diff --git a/deepmd/dpmodel/array_api.py b/deepmd/dpmodel/array_api.py index 115242edfb..ca597b3444 100644 --- a/deepmd/dpmodel/array_api.py +++ b/deepmd/dpmodel/array_api.py @@ -334,6 +334,48 @@ def xp_setitem_at(x: Array, mask: Array, values: Array) -> Array: return x +def xp_uniform(like: Array, size: int, low: float = 0.0, high: float = 1.0) -> Array: + """Draw ``size`` uniform samples in ``[low, high)`` on ``like``'s device. + + Each backend uses its own generator: torch draws with ``torch.rand`` (as + pt does, so ``setup_seed`` replays it, with no host copy -- and a host + draw would freeze to a constant under tracing); other backends use + :mod:`deepmd.utils.random`, which ``setup_seed`` also seeds. Draws are + therefore not comparable across backends -- use only for a per-forward + random stream, never where a parity test looks. + + Parameters + ---------- + like : Array + Reference array supplying backend, dtype and device. + size : int + Number of samples to draw. + low : float + Lower bound of the interval. + high : float + Upper bound of the interval, exclusive. + + Returns + ------- + Array + Samples of shape ``(size,)`` matching ``like``. + """ + if array_api_compat.is_torch_array(like): + import torch + + return ( + torch.rand(size, dtype=like.dtype, device=like.device) * (high - low) + low + ) + from deepmd.utils import random as dp_random + + xp = array_api_compat.array_namespace(like) + drawn = np.asarray(dp_random.random(size)) * (high - low) + low + return xp.astype( + xp_asarray_nodetach(xp, drawn, device=array_api_compat.device(like)), + like.dtype, + ) + + def xp_bincount(x: Array, weights: Array | None = None, minlength: int = 0) -> Array: """Counts the number of occurrences of each value in x.""" xp = array_api_compat.array_namespace(x) diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 36c9fa65e5..a894eac142 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -167,6 +167,29 @@ def has_default_fparam(self) -> bool: """Check if the model has default frame parameters.""" return False + def uses_graph_lower(self) -> bool: + """Returns whether this atomic model supports the NeighborGraph lower. + + Generic capability (concrete default ``False``): the model layer + consults it for graph-route eligibility without assuming anything + about the atomic model's internal architecture. Implementations + answer from their own structure (e.g. a descriptor+fitting model + delegates to its descriptor; a composition supports it iff ALL its + children do). + """ + return False + + def supports_native_spin(self) -> bool: + """Returns whether this atomic model consumes a per-atom spin input. + + Generic capability (concrete default ``False``), the twin of + :meth:`uses_graph_lower`: the model layer asks the atomic model + directly instead of reaching into it for a descriptor, so the answer + stays correct for architectures with no descriptor at all (analytical + terms) and for compositions. + """ + return False + def get_default_fparam(self) -> list[float] | None: """Get the default frame parameters.""" return None @@ -381,6 +404,7 @@ def forward_common_atomic_graph( fparam: Array | None = None, aparam: Array | None = None, charge_spin: Array | None = None, + spin: Array | None = None, comm_dict: dict | None = None, ) -> dict: """Graph analogue of :meth:`forward_common_atomic` on the flat node axis. @@ -405,8 +429,14 @@ def forward_common_atomic_graph( aparam atomic parameter. N x nda charge_spin - charge/spin conditioning. Unused by the dpa1 graph path; accepted so - the interface stays stable for charge/spin-conditioned descriptors. + frame-level charge/spin conditioning, forwarded unchanged to + :meth:`forward_atomic_graph`, which only passes it on to the + descriptor's ``call_graph`` for descriptors that declare + ``supports_charge_spin`` (currently DPA4 only). + spin + flat (N, 3) per-node spin, forwarded unchanged to + :meth:`forward_atomic_graph` (and, from there, the descriptor's + ``call_graph``); None for spin-less models. comm_dict MPI communication metadata forwarded to :meth:`forward_atomic_graph` (and, from there, the descriptor's ``call_graph``). ``None`` for @@ -426,6 +456,7 @@ def forward_common_atomic_graph( fparam=fparam, aparam=aparam, charge_spin=charge_spin, + spin=spin, comm_dict=comm_dict, ) return self._finalize_atomic_ret(ret_dict, output_mask, atype) diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index 9f8f06f82d..1752aabef9 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -123,6 +123,22 @@ def __init__( self.add_chg_spin_ebd: bool = getattr( self.descriptor, "add_chg_spin_ebd", False ) + # Structural capability: only descriptors with a native spin + # conditioning mechanism (currently DPA4) accept a ``spin`` kwarg on + # ``call_graph`` at all -- unlike ``charge_spin``, which every + # descriptor's dense ``call()`` accepts (and ignores) for interface + # stability, the graph-native ``call_graph`` signature is + # per-descriptor, so ``forward_atomic_graph`` must not pass the + # keyword to a descriptor whose ``call_graph`` does not declare it + # (that would be a ``TypeError``, not a no-op). Queried via the + # ``supports_native_spin`` capability method declared on + # ``BaseDescriptor`` (concrete default ``False``; DPA4 overrides). + self._supports_native_spin: bool = self.descriptor.supports_native_spin() + # Same capability method as ``supports_native_spin`` above, for the + # frame-level ``charge_spin`` FiLM kwarg: only DPA4's ``call_graph`` + # declares it; other descriptors' ``call_graph`` would ``TypeError`` + # on an unconditional ``charge_spin=`` kwarg. + self.supports_charge_spin: bool = self.descriptor.supports_charge_spin() super().init_out_stat() def has_chg_spin_ebd(self) -> bool: @@ -147,6 +163,14 @@ def get_default_chg_spin(self) -> list[float] | None: return self.descriptor.get_default_chg_spin() return None + def uses_graph_lower(self) -> bool: + """Delegates to this model's own descriptor.""" + return bool(self.descriptor.uses_graph_lower()) + + def supports_native_spin(self) -> bool: + """Delegates to this model's own descriptor (cached at construction).""" + return self._supports_native_spin + def fitting_output_def(self) -> FittingOutputDef: """Get the output def of the fitting net.""" return self.fitting_net.output_def() @@ -301,6 +325,7 @@ def forward_atomic_graph( fparam: Array | None = None, aparam: Array | None = None, charge_spin: Array | None = None, + spin: Array | None = None, comm_dict: dict | None = None, ) -> dict[str, Array]: """Graph analogue of :meth:`forward_atomic` on the flat node axis. @@ -321,8 +346,14 @@ def forward_atomic_graph( aparam atomic parameter. N x nda charge_spin - charge/spin conditioning. Unused by the dpa1 graph path; accepted so - the interface stays stable for charge/spin-conditioned descriptors. + frame-level charge/spin conditioning, forwarded to the + descriptor's ``call_graph`` only when + ``self.supports_charge_spin`` (currently DPA4 only); ignored (not + forwarded, never a ``TypeError``) for descriptors without that + capability, keeping the interface stable for all of them. + spin + flat (N, 3) per-node spin, forwarded to the descriptor's + ``call_graph``; None for spin-less models. comm_dict MPI communication metadata forwarded to the descriptor's ``call_graph`` (the message-passing part). ``None`` for @@ -342,9 +373,24 @@ def forward_atomic_graph( ) xp = array_api_compat.array_namespace(graph.edge_vec) - type_embedding = self.descriptor.type_embedding.call() + # Descriptor-owned: dpa1/dpa2 hand out their full tebd table; DPA4 + # embeds types internally from ``atype`` and returns None. + type_embedding = self.descriptor.graph_type_embedding_table() + # Only forward the ``spin``/``charge_spin`` keyword to descriptors + # whose ``call_graph`` declares it. Queried through the public + # capability -- an override must reach THIS call site too, so the + # cached field stays behind the method. + spin_kwargs = {"spin": spin} if self.supports_native_spin() else {} + charge_spin_kwargs = ( + {"charge_spin": charge_spin} if self.supports_charge_spin else {} + ) gg, rot_mat = self.descriptor.call_graph( - graph, atype, type_embedding=type_embedding, comm_dict=comm_dict + graph, + atype, + type_embedding=type_embedding, + comm_dict=comm_dict, + **spin_kwargs, + **charge_spin_kwargs, ) fparam_node = None if fparam is not None: diff --git a/deepmd/dpmodel/atomic_model/inter_potential.py b/deepmd/dpmodel/atomic_model/inter_potential.py new file mode 100644 index 0000000000..b95aeb8da5 --- /dev/null +++ b/deepmd/dpmodel/atomic_model/inter_potential.py @@ -0,0 +1,485 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Analytical pair potentials for Zone bridging (backend-agnostic port of +``deepmd.pt``'s ``InterPotential``). Lives in the atomic-model package: +the atomic layer owns per-atom energy assembly, where the ZBL term is +injected on the graph route. +""" + +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel.array_api import ( + Array, + xp_asarray_nodetach, +) +from deepmd.dpmodel.common import ( + NativeOP, +) +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + OutputVariableDef, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .base_atomic_model import ( + BaseAtomicModel, +) + +# fmt: off +ELEMENT_TO_Z: dict[str, int] = { + "H": 1, "He": 2, "Li": 3, "Be": 4, "B": 5, "C": 6, "N": 7, "O": 8, + "F": 9, "Ne": 10, "Na": 11, "Mg": 12, "Al": 13, "Si": 14, "P": 15, + "S": 16, "Cl": 17, "Ar": 18, "K": 19, "Ca": 20, "Sc": 21, "Ti": 22, + "V": 23, "Cr": 24, "Mn": 25, "Fe": 26, "Co": 27, "Ni": 28, "Cu": 29, + "Zn": 30, "Ga": 31, "Ge": 32, "As": 33, "Se": 34, "Br": 35, "Kr": 36, + "Rb": 37, "Sr": 38, "Y": 39, "Zr": 40, "Nb": 41, "Mo": 42, "Tc": 43, + "Ru": 44, "Rh": 45, "Pd": 46, "Ag": 47, "Cd": 48, "In": 49, "Sn": 50, + "Sb": 51, "Te": 52, "I": 53, "Xe": 54, "Cs": 55, "Ba": 56, "La": 57, + "Ce": 58, "Pr": 59, "Nd": 60, "Pm": 61, "Sm": 62, "Eu": 63, "Gd": 64, + "Tb": 65, "Dy": 66, "Ho": 67, "Er": 68, "Tm": 69, "Yb": 70, "Lu": 71, + "Hf": 72, "Ta": 73, "W": 74, "Re": 75, "Os": 76, "Ir": 77, "Pt": 78, + "Au": 79, "Hg": 80, "Tl": 81, "Pb": 82, "Bi": 83, "Po": 84, "At": 85, + "Rn": 86, "Fr": 87, "Ra": 88, "Ac": 89, "Th": 90, "Pa": 91, "U": 92, + "Np": 93, "Pu": 94, "Am": 95, "Cm": 96, "Bk": 97, "Cf": 98, "Es": 99, + "Fm": 100, "Md": 101, "No": 102, "Lr": 103, "Rf": 104, "Db": 105, + "Sg": 106, "Bh": 107, "Hs": 108, "Mt": 109, "Ds": 110, "Rg": 111, + "Cn": 112, "Nh": 113, "Fl": 114, "Mc": 115, "Lv": 116, "Ts": 117, + "Og": 118, +} +# fmt: on + +# ZBL screening function coefficients +_ZBL_A_COEFF = (0.18175, 0.50986, 0.28022, 0.028171) +_ZBL_B_COEFF = (3.1998, 0.94229, 0.4029, 0.20162) + +# Physical constants +_KE_EV_A = 14.3996 # Coulomb constant in eV·Å +_A_BOHR = 0.5291772109 # Bohr radius in Å + + +class InterPotential(NativeOP): + """Analytical pair potential for Zone bridging. + + Supports the Ziegler-Biersack-Littmark (ZBL) screened nuclear repulsion + potential, evaluated on the edge form so that its force and virial flow + through the same edge backward as the learned energy. Each pair (i, j) + contributes ``V_ZBL(r_ij) / 2`` to both atom i and atom j, avoiding + double-counting from the symmetric neighbor list. Backend-agnostic + (array-API) port of the reference implementation in + ``deepmd.pt.model.model.sezm_model.InterPotential``. + + Parameters + ---------- + type_map : list[str] + Element symbols (e.g. ``["O", "H"]``). Index in this list + corresponds to the ``atype`` integer values. + mode : str + Potential formula. Currently only ``"zbl"`` is supported. + + Raises + ------ + ValueError + If ``mode`` is not recognized, or if any element in ``type_map`` is + not found in the periodic table. + """ + + def __init__(self, type_map: list[str], mode: str = "zbl") -> None: + super().__init__() + mode = str(mode).upper() + if mode != "ZBL": + raise ValueError(f"Unknown InterPotential mode: {mode}") + self.mode = mode + self.type_map = list(type_map) + self.ntypes_real = len(type_map) + self.atomic_numbers = self._lookup_from_type_map(type_map) + + @staticmethod + def _lookup_from_type_map(type_map: list[str], like: Array | None = None) -> Array: + """Build the per-type nuclear-charge lookup from element symbols. + + Parameters + ---------- + type_map : list[str] + Element symbols; index corresponds to ``atype`` values. + like : Array, optional + When given, the result is created in this array's namespace, dtype + and device instead of NumPy -- so an in-place rebuild on a wrapped + backend (pt_expt buffer, possibly on CUDA) stays where it was. + + Returns + ------- + Array + Nuclear charges, shape ``(len(type_map),)``. + + Raises + ------ + ValueError + If an element symbol is not in :data:`ELEMENT_TO_Z`. + """ + atomic_numbers = [] + for elem in type_map: + z = ELEMENT_TO_Z.get(elem) + if z is None: + raise ValueError(f"Unknown element symbol: {elem}") + atomic_numbers.append(z) + arr = np.asarray(atomic_numbers, dtype=np.float64) + if like is None: + return arr + xp = array_api_compat.array_namespace(like) + return xp.asarray(arr, dtype=like.dtype, device=array_api_compat.device(like)) + + def change_type_map(self, type_map: list[str]) -> None: + """Rebuild the element lookup for a new type map. + + THIS OWNS the element lookup, so it owns every update of it: the + symbols, their count and the nuclear-charge table are one piece of + state and are replaced together. Reordering, adding and dropping + elements are all covered -- the table is rebuilt from the symbols + rather than permuted, so no index bookkeeping can drift. The rebuilt + array keeps the current one's namespace/dtype/device, so a wrapped + backend (pt_expt buffer on CPU or CUDA) is updated in place. + + Parameters + ---------- + type_map : list[str] + The new element symbols. + """ + self.atomic_numbers = self._lookup_from_type_map( + type_map, like=self.atomic_numbers + ) + self.type_map = list(type_map) + self.ntypes_real = len(type_map) + + @staticmethod + def _zbl_pair_energy(xp: Any, r: Array, zi: Array, zj: Array) -> Array: + """Compute ZBL pair energy for given distances and nuclear charges. + + Parameters + ---------- + xp + The array namespace of ``r``/``zi``/``zj``. + r : Array + Pair distances with shape (...) in Å. + zi : Array + Nuclear charge of atom i with shape (...). + zj : Array + Nuclear charge of atom j with shape (...). + + Returns + ------- + Array + Pair energies with shape (...) in eV. + """ + a_screen = 0.88534 * _A_BOHR / (zi**0.23 + zj**0.23) + x = r / a_screen + phi = sum( + a_k * xp.exp(-b_k * x) + for a_k, b_k in zip(_ZBL_A_COEFF, _ZBL_B_COEFF, strict=True) + ) + return _KE_EV_A * zi * zj / r * phi + + def call( + self, + edge_vec: Array, + edge_index: Array, + atype_flat: Array, + edge_mask: Array, + n_node: int, + real_type_count: int | None = None, + ) -> Array: + """Scatter per-edge ZBL half-energies into per-atom energies. + + Parameters + ---------- + edge_vec : Array + (E, 3) edge vectors in Å (the autograd leaf on differentiable + backends: differentiating the returned energy w.r.t. this input + yields the ZBL force/virial through the shared edge backward). + edge_index : Array + (2, E) ``[src, dst]`` edge endpoints (flat node indices). + atype_flat : Array + (N,) flat atom types. + edge_mask : Array + (E,) valid-edge mask. + n_node : int + Total flat node count ``N``. + real_type_count : int | None + Count of REAL atom types; types ``>= real_type_count`` + (virtual/placeholder) are wrapped back to their real parent for + the Z lookup and masked out of the sum. Defaults to + ``len(type_map)``. + + Returns + ------- + Array + Per-atom ZBL energies with shape ``(1, n_node, 1)`` in + ``edge_vec``'s dtype. + """ + xp = array_api_compat.array_namespace(edge_vec) + device = array_api_compat.device(edge_vec) + if real_type_count is None: + real_type_count = self.ntypes_real + src = xp.astype(edge_index[0, :], xp.int64) + dst = xp.astype(edge_index[1, :], xp.int64) + r = xp.linalg.vector_norm(xp.astype(edge_vec, xp.float64), axis=-1) + r = xp.clip(r, min=1e-10) + # Virtual/placeholder types wrap back to the real parent purely so + # the Z lookup never indexes out of range; their edges are masked + # out below. + atype_i64 = xp.astype(atype_flat, xp.int64) + atype_for_z = xp.clip(atype_i64, min=0) + atype_for_z = xp.where( + atype_for_z >= real_type_count, + atype_for_z - real_type_count, + atype_for_z, + ) + z_all = xp.take( + xp_asarray_nodetach( + xp, self.atomic_numbers, dtype=xp.float64, device=device + ), + atype_for_z, + axis=0, + ) + zi = xp.take(z_all, src, axis=0) + zj = xp.take(z_all, dst, axis=0) + pair_e = self._zbl_pair_energy(xp, r, zi, zj) + node_is_real = atype_i64 < real_type_count + valid = ( + xp.astype(edge_mask, xp.bool) + & xp.take(node_is_real, src, axis=0) + & xp.take(node_is_real, dst, axis=0) + ) + pair_e = pair_e * xp.astype(valid, pair_e.dtype) + # Symmetric neighbor list: both directed edges exist, each scatters + # half into its dst -- atoms i and j each receive V/2. + from deepmd.dpmodel.utils.neighbor_graph import ( + segment_sum, + ) + + atom_energy = segment_sum(pair_e * 0.5, dst, n_node) + return xp.astype(xp.reshape(atom_energy, (1, n_node, 1)), edge_vec.dtype) + + +@BaseAtomicModel.register("inter_potential") +class InterPotentialAtomicModel(BaseAtomicModel): + """Analytical bridging pair potential as an ATOMIC MODEL. + + First-principles composition design: the analytical term maps local + atomic environments to per-atom energies -- exactly the atomic-model + contract -- so a "bridging model" is a SUM of two atomic energy models + (the learned one and this one) via + :class:`~deepmd.dpmodel.atomic_model.linear_atomic_model.LinearEnergyAtomicModel` + with ``weights="sum"``, not a flag on the learned model. Graph-route + only: the term is evaluated on the shared ``graph.edge_vec`` leaf so + its force/virial ride the same edge backward as the learned energy; + the dense (nlist) route raises. + + Parameters + ---------- + type_map : list[str] + Element symbols; index corresponds to ``atype`` values. + mode : str + Potential formula (currently ``"zbl"``). + rcut : float + Cut-off radius this model declares (the composition uses the max + over children; pass the learned model's). + sel : list[int] | int + Neighbor selection this model declares (composition bookkeeping). + """ + + def __init__( + self, + type_map: list[str], + mode: str = "zbl", + rcut: float = 0.0, + sel: "list[int] | int" = 0, + **kwargs: Any, + ) -> None: + super().__init__(type_map, **kwargs) + self.potential = InterPotential(type_map=list(type_map), mode=mode) + self.mode = self.potential.mode + self.rcut = float(rcut) + self.sel = ( + [int(s) for s in sel] if isinstance(sel, (list, tuple)) else [int(sel)] + ) + super().init_out_stat() + + def change_type_map( + self, type_map: list[str], model_with_new_type_stat: Any | None = None + ) -> None: + """Change the type related params to new ones, according to `type_map` and the original one in the model. + If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. + + The generic base handles the public map and the stat/exclusion state; + the element lookup belongs to :class:`InterPotential`, so the update is + delegated there rather than reimplemented here (review 3649295675 -- + without it the lookup keeps the ORIGINAL elements while ``atype`` + values mean new ones, and a longer new map raises ``IndexError``). + + Parameters + ---------- + type_map : list[str] + The new element symbols. + model_with_new_type_stat : optional + Model with statistics for the new types (unused: an analytical + term has no fitted statistics). + """ + super().change_type_map( + type_map, model_with_new_type_stat=model_with_new_type_stat + ) + self.potential.change_type_map(type_map) + + def fitting_output_def(self) -> FittingOutputDef: + """Per-atom analytical energy: reducible and fully differentiable.""" + return FittingOutputDef( + [ + OutputVariableDef( + name="energy", + shape=[1], + reducible=True, + r_differentiable=True, + c_differentiable=True, + ) + ] + ) + + def get_rcut(self) -> float: + """Get the cut-off radius.""" + return self.rcut + + def get_sel(self) -> list[int]: + """Get the neighbor selection.""" + return self.sel + + def get_nsel(self) -> int: + """Get the total neighbor selection.""" + return sum(self.sel) + + def mixed_types(self) -> bool: + """The analytical term is type-agnostic in layout (mixed types).""" + return True + + def has_message_passing(self) -> bool: + """No message passing in an analytical pair term.""" + return False + + def need_sorted_nlist_for_lower(self) -> bool: + """No nlist ordering requirement (graph-route only).""" + return False + + def get_dim_fparam(self) -> int: + """No frame parameters.""" + return 0 + + def get_dim_aparam(self) -> int: + """No atomic parameters.""" + return 0 + + def get_sel_type(self) -> list[int]: + """All atom types contribute.""" + return [] + + def is_aparam_nall(self) -> bool: + """No atomic parameters.""" + return False + + def uses_graph_lower(self) -> bool: + """Graph-only term: the NeighborGraph lower is its sole evaluation + route, so it supports the graph route like any graph-capable atomic + model. + """ + return True + + def forward_atomic( + self, + *args: Any, + **kwargs: Any, + ) -> dict: + """Dense route unsupported: the term rides the NeighborGraph route only.""" + raise NotImplementedError( + "InterPotentialAtomicModel rides the NeighborGraph route only; " + "the dense (nlist) route has no injection site for the term" + ) + + def forward_atomic_graph( + self, + graph: Any, + atype: Any, + fparam: Any = None, + aparam: Any = None, + charge_spin: Any = None, + spin: Any = None, + comm_dict: dict | None = None, + ) -> dict: + """Evaluate the analytical per-atom energy on the flat node axis. + + ``fparam``/``aparam``/``charge_spin``/``spin``/``comm_dict`` are + accepted for pipeline-signature compatibility and ignored (the + analytical term conditions on geometry and types only). + + Parameters + ---------- + graph + neighbor graph; ``graph.edge_vec`` is the differentiable edge + leaf on autograd backends. + atype + flat local atom types. N + + Returns + ------- + dict + ``{"energy": (N, 1)}`` per-atom analytical energies. + """ + import array_api_compat + + xp = array_api_compat.array_namespace(graph.edge_vec) + n_node = atype.shape[0] + energy = self.potential.call( + graph.edge_vec, + graph.edge_index, + atype, + graph.edge_mask, + n_node=n_node, + real_type_count=len(self.type_map), + ) + return {"energy": xp.reshape(energy, (n_node, 1))} + + def serialize(self) -> dict: + data = super().serialize() + data.update( + { + "@class": "Model", + "type": "inter_potential", + "@version": 1, + "mode": self.mode, + "rcut": self.rcut, + "sel": self.sel, + } + ) + return data + + @classmethod + def deserialize(cls, data: dict) -> "InterPotentialAtomicModel": + data = data.copy() + check_version_compatibility(data.pop("@version", 1), 1, 1) + data.pop("@class", None) + data.pop("type", None) + return super().deserialize(data) + + def set_case_embd(self, case_idx: int) -> None: + """No case embedding in an analytical term.""" + + def compute_or_load_stat( + self, + sampled_func: Any, + stat_file_path: Any = None, + compute_or_load_out_stat: bool = True, + preset_observed_type: "list[str] | None" = None, + ) -> None: + """Analytical term: no statistics to compute.""" diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 5e70f5c0ed..293ad005b5 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -40,6 +40,7 @@ @BaseAtomicModel.register("linear") +@BaseAtomicModel.register("linear_ener") # accepted alias, never emitted class LinearEnergyAtomicModel(BaseAtomicModel): r"""Linear model makes linear combinations of several existing models. @@ -83,9 +84,37 @@ def __init__( f"LinearAtomicModel only supports AtomicModel of mixed type, the following models are not mixed type: {model_mixed_type}." ) + # Fail fast: a sum mixing an intensive with an extensive term is not + # physically meaningful, so such a composition must not exist. + intensive_flags = {m.get_intensive() for m in models} + if len(intensive_flags) > 1: + raise ValueError( + "LinearAtomicModel cannot combine intensive and extensive " + "sub-models: " + + ", ".join(f"{type(m).__name__}={m.get_intensive()}" for m in models) + ) + + # Fail fast: the composition feeds ONE external fparam/aparam tensor to + # every child, so all children that actually consume it must agree on + # its dimension. Children with dimension 0 do not consume it and are + # ignored, so a learned model composed with an analytical term (ZBL) + # simply inherits the learned dimension. + for name, dim_of in ( + ("fparam", lambda m: m.get_dim_fparam()), + ("aparam", lambda m: m.get_dim_aparam()), + ): + dims = {dim_of(m) for m in models if dim_of(m) > 0} + if len(dims) > 1: + raise ValueError( + f"LinearAtomicModel sub-models disagree on the {name} " + "dimension, but the composition feeds them one shared " + "tensor: " + + ", ".join(f"{type(m).__name__}={dim_of(m)}" for m in models) + ) + self.models = models self.type_map = type_map - self.mapping_list = self._build_mapping_list() + self._rebuild_mapping_state() self.mixed_types_list = [model.mixed_types() for model in self.models] if isinstance(weights, str): assert weights in ["sum", "mean"] @@ -97,6 +126,21 @@ def __init__( ) self.weights = weights + def _rebuild_mapping_state(self) -> None: + """Rebuild ``mapping_list`` and everything derived from it. + + The ONE owner of the mapping state: called at construction and after + ``change_type_map`` (submodels may reorder or add species). + ``_graph_mapping_is_identity`` is a static composition property + computed EAGERLY here: the graph route requires identity atype + mappings, and checking at forward time would iterate (possibly + traced) tensors and trip ``torch.export``'s data-dependent guards. + """ + self.mapping_list = self._build_mapping_list() + self._graph_mapping_is_identity = all( + list(m) == list(range(len(m))) for m in self.mapping_list + ) + def _build_mapping_list(self) -> list[Array]: """Map common type IDs to the current type IDs of every submodel.""" common_type_map = set(self.type_map) @@ -167,8 +211,9 @@ def change_type_map( else None, ) # Submodels may reorder existing species or add new ones. Rebuild only - # after every submodel has changed so runtime type IDs use their new maps. - self.mapping_list = self._build_mapping_list() + # after every submodel has changed so runtime type IDs use their new maps + # (also refreshes the derived _graph_mapping_is_identity flag). + self._rebuild_mapping_state() def get_model_rcuts(self) -> list[float]: """Get the cut-off radius for each individual models.""" @@ -233,6 +278,116 @@ def enable_compression( check_frequency, ) + def uses_graph_lower(self) -> bool: + """Graph-capable iff EVERY child supports the graph lower. + + All children evaluate on the one shared graph, so a single + dense-only child (e.g. ``PairTabAtomicModel`` in standard DP+ZBL) + forces the whole composition onto the dense route; a graph + descriptor plus an analytical graph term (ZBL bridging) stays on the + graph route. + """ + return all(m.uses_graph_lower() for m in self.models) + + def supports_native_spin(self) -> bool: + """Spin-capable when ANY child consumes the spin input. + + Unlike :meth:`uses_graph_lower` (every child must run on the shared + graph), spin only has to reach ONE consumer: analytical terms accept + and ignore it, so a composition of a spin-aware learned model with a + ZBL term is a valid native-spin model. With no consumer at all the + magnetic force would be identically zero, which is not a spin model. + """ + return any(m.supports_native_spin() for m in self.models) + + def forward_atomic_graph( + self, + graph: Any, + atype: Array, + fparam: Array | None = None, + aparam: Array | None = None, + charge_spin: Array | None = None, + spin: Array | None = None, + comm_dict: dict | None = None, + ) -> dict[str, Array]: + """Graph-route linear combination on the flat node axis. + + Every child consumes the SAME graph, so on autograd backends the + shared ``graph.edge_vec`` leaf makes the summed energy's force and + virial exactly the sum of the children's -- one edge backward + covers the whole composition (this is what makes analytical + bridging terms compose with the learned model for free). + + Only constant weights are supported here (``"sum"``/``"mean"`` or a + per-child float list); the distance-switched ZBL-interpolation + weights are a dense-route feature. Children with distinct type maps + are not supported on the graph route. + + Parameters + ---------- + graph + neighbor graph for the local atoms (ghost-free). + atype + flat local atom types. N + fparam + frame parameter. nf x ndf + aparam + atomic parameter. N x nda + charge_spin + frame-level conditioning, forwarded to every child (children + gate it on their own capabilities). + spin + flat (N, 3) per-node spin, forwarded to every child. + comm_dict + MPI communication metadata, forwarded to every child. + + Returns + ------- + dict + ``{"energy": (N, 1)}`` -- the weighted sum of the children's + per-atom energies. + + Raises + ------ + NotImplementedError + For non-constant weights or children with remapped type maps. + """ + import array_api_compat + + if not self._graph_mapping_is_identity: + raise NotImplementedError( + "the graph route supports children sharing the parent " + "type_map only (no atype remapping)" + ) + nmodels = len(self.models) + if self.weights == "sum": + weights = [1.0] * nmodels + elif self.weights == "mean": + weights = [1.0 / nmodels] * nmodels + elif isinstance(self.weights, list): + weights = [float(w) for w in self.weights] + else: + raise NotImplementedError( + "the graph route supports constant weights only " + "('sum'/'mean'/list); distance-switched weights are a " + "dense-route feature" + ) + xp = array_api_compat.array_namespace(graph.edge_vec) + energy = None + for model, ww in zip(self.models, weights, strict=True): + ret = model.forward_common_atomic_graph( + graph, + atype, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + spin=spin, + comm_dict=comm_dict, + ) + contrib = ret["energy"] * ww + energy = contrib if energy is None else energy + contrib + return {"energy": xp.astype(energy, graph.edge_vec.dtype)} + def forward_atomic( self, extended_coord: Array, @@ -356,6 +511,10 @@ def serialize(self) -> dict: { "@class": "Model", "@version": 3, + # ONE wire type across backends: pt/tf write "linear" here, so + # dpmodel must too or cross-backend conversion breaks. The + # unambiguous energy-specific name lives in the config/model + # registry ("linear_ener"), which is also accepted here. "type": "linear", "models": [model.serialize() for model in self.models], "type_map": self.type_map, @@ -460,14 +619,117 @@ def _compute_weight( raise NotImplementedError def get_dim_fparam(self) -> int: - """Get the number (dimension) of frame parameters of this atomic model.""" - # tricky... + """Get the number (dimension) of frame parameters of this atomic model. + + ``max`` is exact here, not a guess: ``__init__`` rejects consumers + that disagree, so the only other value present is 0 from a + non-consumer. + """ return max([model.get_dim_fparam() for model in self.models]) def get_dim_aparam(self) -> int: - """Get the number (dimension) of atomic parameters of this atomic model.""" + """Get the number (dimension) of atomic parameters of this atomic model. + + ``max`` is exact for the same reason as :meth:`get_dim_fparam`. + """ return max([model.get_dim_aparam() for model in self.models]) + # --- conditioning-input capabilities owned by the children ----------- + # A composition must FORWARD every capability its children own, exactly + # like get_dim_fparam/get_dim_aparam above. Falling through to + # BaseAtomicModel's False/0 is silently wrong: eager forward still + # conditions on the input (the learned child consumes it), but the + # freeze reads these accessors, so a 0 here drops the charge_spin slot + # from the exported ABI and from the metadata the C++ feeder uses -- + # the artifact then disagrees with its own eager model. + + def has_chg_spin_ebd(self) -> bool: + """Whether ANY child consumes the frame-level charge/spin FiLM input.""" + return any(model.has_chg_spin_ebd() for model in self.models) + + def get_intensive(self) -> bool: + """Whether the composed property is intensive. + + All children agree by construction (validated in ``__init__``). + """ + return self.models[0].get_intensive() if self.models else False + + def get_compute_stats_distinguish_types(self) -> bool: + """Needed if ANY child needs them; the stricter rule is safe for + children that do not distinguish types. + """ + return any(model.get_compute_stats_distinguish_types() for model in self.models) + + def get_dim_chg_spin(self) -> int: + """Dimension of the charge_spin input (max over children, like fparam).""" + return max([model.get_dim_chg_spin() for model in self.models]) + + @staticmethod + def _agreed_default( + actives: "list[BaseAtomicModel]", + has: "Callable[[BaseAtomicModel], bool]", + get: "Callable[[BaseAtomicModel], Any]", + ) -> "tuple[bool, Any]": + """Shared default of the ACTIVE children, or none if they disagree. + + The composition exposes ONE external tensor to all children, so a + parent default is only meaningful when every active consumer would + have used the same value anyway. Otherwise omitting the input must + stay omitted, letting each child apply its own default, rather than + silently broadcasting one child's value to the others. + + Children that do not consume the input (dimension 0, e.g. an + analytical bridging term) are excluded by the caller, so a learned + model composed with ZBL still inherits the learned default. + """ + if not actives or not all(has(m) for m in actives): + return False, None + values = [np.asarray(get(m), dtype=float).reshape(-1) for m in actives] + first = values[0] + if any(v.shape != first.shape or not np.allclose(v, first) for v in values[1:]): + return False, None + return True, get(actives[0]) + + def _chg_spin_consumers(self) -> list: + """Children that actually consume ``charge_spin``.""" + return [m for m in self.models if m.get_dim_chg_spin() > 0] + + def _fparam_consumers(self) -> list: + """Children that actually consume ``fparam``.""" + return [m for m in self.models if m.get_dim_fparam() > 0] + + def has_default_chg_spin(self) -> bool: + """Whether every active child shares one default charge/spin.""" + return self._agreed_default( + self._chg_spin_consumers(), + lambda m: m.has_default_chg_spin(), + lambda m: m.get_default_chg_spin(), + )[0] + + def get_default_chg_spin(self) -> "Array | None": + """The shared default charge/spin conditions, if the children agree.""" + return self._agreed_default( + self._chg_spin_consumers(), + lambda m: m.has_default_chg_spin(), + lambda m: m.get_default_chg_spin(), + )[1] + + def has_default_fparam(self) -> bool: + """Whether every active child shares one default frame parameter.""" + return self._agreed_default( + self._fparam_consumers(), + lambda m: m.has_default_fparam(), + lambda m: m.get_default_fparam(), + )[0] + + def get_default_fparam(self) -> "list[float] | None": + """The shared default frame parameters, if the children agree.""" + return self._agreed_default( + self._fparam_consumers(), + lambda m: m.has_default_fparam(), + lambda m: m.get_default_fparam(), + )[1] + def get_sel_type(self) -> list[int]: """Get the selected atom types of this model. diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 24e85ae678..eb8eeccbf6 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -459,6 +459,16 @@ def uses_graph_lower(self) -> bool: ) return self.se_atten.tebd_input_mode in ("concat", "strip") + def graph_type_embedding_table(self) -> Array: + """Full type-embedding table consumed by the graph-route forward. + + Returns + ------- + Array + The ``(ntypes + 1, tebd_dim)`` table from ``type_embedding``. + """ + return self.type_embedding.call() + def uses_compact_edge_pairs(self) -> bool: """Returns whether the graph lower traces compact edge pairs. diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index 06a86956ee..ccbfba085d 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -751,6 +751,16 @@ def uses_graph_lower(self) -> bool: return False return self.repinit.tebd_input_mode in ("concat", "strip") + def graph_type_embedding_table(self) -> Array: + """Full type-embedding table consumed by the graph-route forward. + + Returns + ------- + Array + The ``(ntypes + 1, tebd_dim)`` table from ``type_embedding``. + """ + return self.type_embedding.call() + def uses_compact_edge_pairs(self) -> bool: """Returns whether the graph lower traces compact edge pairs. diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index 11f96c4669..2687d9535d 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -50,7 +50,6 @@ from deepmd.dpmodel.array_api import ( xp_asarray_nodetach, xp_scatter_sum, - xp_take_first_n, ) from deepmd.dpmodel.common import ( PRECISION_DICT, @@ -63,6 +62,10 @@ from deepmd.dpmodel.utils.exclude_mask import ( PairExcludeMask, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + apply_pair_exclusion, + graph_from_dense_quartet, +) from deepmd.dpmodel.utils.seed import ( child_seed, ) @@ -84,8 +87,7 @@ ) from .dpa4_nn.edge_cache import ( EdgeCache, - build_edge_cache, - build_edge_cache_from_edges, + _edge_cache_from_arrays, edge_cache_to_dtype, ) from .dpa4_nn.embedding import ( @@ -129,6 +131,9 @@ from deepmd.dpmodel.array_api import ( Array, ) + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + ) from deepmd.utils.data_system import ( DeepmdDataSystem, ) @@ -137,6 +142,86 @@ ) +def _graph_from_padded_nlist( + coord_ext: Array, + atype_ext: Array, + nlist: Array, + mapping: Array | None, +) -> tuple[NeighborGraph, Array]: + """Dense-topology extraction for DPA4: ``src_ok``-sanitized nlist -> graph. + + This is the ONE owner of DPA4's dense-nlist topology contract. It + reproduces the pt SeZM ``edge_keep`` semantics (pt's ``src_ok`` filter, + see ``deepmd.pt.model.descriptor.sezm_nn.edge_cache.build_edge_cache``): + beyond the native ``nlist == -1`` padding, a neighbor slot is also + invalid when its LOCAL source index falls outside ``[0, nloc)`` -- + + - ``mapping`` given: the mapped owner ``mapping[nlist]`` is out of range + (covers broken ghosts with ``mapping == -1``, which pt drops and the + retired dp dense builder masked); + - ``mapping is None`` (neighbor indices already local): the entry itself + is out of ``[0, nloc)``. + + Invalid slots are rewritten to ``-1`` BEFORE conversion, so + :func:`graph_from_dense_quartet` treats them as native padding (masked + edge, in-range placeholder indices, zero geometry) -- semantically + identical to the retired dense builder's masked slots. The gather that + tests the mapped source uses ``where(nlist >= 0, nlist, 0)`` clamping, + so no negative index is ever gathered. The ``mapping is None`` identity + case itself is owned by :func:`graph_from_dense_quartet` (applied after + this sanitization). This contract is pinned by + ``source/tests/pt/model/test_dpa4_dpmodel_parity.py`` + (``TestEdgeCacheParity`` and the descriptor-parity fixtures with a + broken ghost mapping). + + Parameters + ---------- + coord_ext + Extended coordinates with shape (nf, nall, 3), already in the + caller's compute precision (edge vectors are subtracted inside the + converter in this dtype). + atype_ext + Extended atom types with shape (nf, nall). + nlist + Padded neighbor list with shape (nf, nloc, nnei); ``-1`` is padding. + mapping + Extended -> local-owner index with shape (nf, nall), or ``None`` + when neighbor indices are already local. + + Returns + ------- + tuple[NeighborGraph, Array] + The shape-static row-major graph over the local atoms and the flat + local atom types with shape (nf * nloc,). + """ + xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist) + device = array_api_compat.device(nlist) + nf, nloc, _ = nlist.shape + nlist = xp.astype(nlist, xp.int64) + valid = nlist >= 0 + nl_safe = xp.where(valid, nlist, xp.zeros_like(nlist)) + if mapping is None: + # Neighbor indices are already local owners. + src_local = nl_safe + else: + nall = atype_ext.shape[1] + mapping_flat = xp.astype(xp.reshape(mapping, (-1,)), xp.int64) + frame_idx = xp.reshape(xp.arange(nf, dtype=xp.int64, device=device), (nf, 1, 1)) + src_local = xp.reshape( + xp.take( + mapping_flat, + xp.reshape(frame_idx * nall + nl_safe, (-1,)), + axis=0, + ), + nlist.shape, + ) + valid = valid & (src_local >= 0) & (src_local < nloc) + nlist_sane = xp.where( + valid, nl_safe, xp.full(nlist.shape, -1, dtype=xp.int64, device=device) + ) + return graph_from_dense_quartet(coord_ext, atype_ext, nlist_sane, mapping) + + @BaseDescriptor.register("SeZM") @BaseDescriptor.register("sezm") @BaseDescriptor.register("DPA4") @@ -202,7 +287,7 @@ class DescrptDPA4(NativeOP, BaseDescriptor): Execution outline ----------------- - 1. Build a per-forward `EdgeFeatureCache` (geometry, envelope, Wigner-D). + 1. Build a per-forward `EdgeCache` (geometry, envelope, Wigner-D). 2. Build radial/type edge features once and reuse across blocks. 3. Run `SeZMInteractionBlock` stack with optional l/m schedules. 4. Extract scalar channels and apply the final scalar FFN. @@ -580,6 +665,7 @@ def __init__( **kwargs: Any, ) -> None: self.version = float(self.LATEST_VERSION) + self._graph_lower_disabled = False self.rcut = float(rcut) if env_exp is None: env_exp = [7, 5] @@ -1153,9 +1239,6 @@ def call( atype_ext: Array, nlist: Array, mapping: Array | None = None, - edge_index: Array | None = None, - edge_vec: Array | None = None, - edge_mask: Array | None = None, comm_dict: dict[str, Array] | None = None, fparam: Array | None = None, force_embedding: Array | None = None, @@ -1181,17 +1264,11 @@ def call( Neighbor list with shape (nf, nloc, nnei). mapping Extended-to-local mapping with shape (nf, nall), or None. - edge_index - Fixed-shape edge indices with shape (2, E). If provided, the descriptor - uses the edge-list path and ignores `nlist` and `mapping`. - edge_vec - Fixed-shape edge vectors with shape (E, 3) in Å. Required when - `edge_index` is provided. - edge_mask - Fixed-shape edge mask with shape (E,). Required when `edge_index` - is provided. comm_dict - Communication dictionary for parallel inference (unused). + Communication dictionary for parallel inference. The dense + (nlist) lower does not implement multi-rank exchange; a non-None + value raises ``NotImplementedError`` here, before the graph + build. fparam Frame parameters with shape (nf, nfp). Not used by SeZM, kept for interface compatibility. @@ -1215,6 +1292,11 @@ def call( None (not used). sw None (not used). + + Raises + ------ + NotImplementedError + When ``comm_dict`` is provided. """ xp = array_api_compat.array_namespace(coord_ext, atype_ext) device = array_api_compat.device(coord_ext) @@ -1223,320 +1305,231 @@ def call( elif coord_ext.ndim != 3: raise ValueError("coord_ext must have shape (nf, nall*3) or (nf, nall, 3)") - if edge_index is not None: - nf_edge = atype_ext.shape[0] - charge_spin = self._canonicalize_charge_spin( - charge_spin, - nf=nf_edge, - dtype=coord_ext.dtype, - device=device, - ) - descriptor, _ = self.call_with_edges( - coord_ext=coord_ext, - atype_ext=atype_ext, - edge_index=edge_index, - edge_vec=edge_vec, - edge_mask=edge_mask, - force_embedding=force_embedding, - charge_spin=charge_spin, - spin=spin, - ) - return ( - descriptor, - None, - None, - None, - None, - ) - - # === Step 1. Setup dimensions === + # === Dense-nlist adapter over the graph-native core === + # ``graph_from_dense_quartet`` enumerates edges row-major over + # (frame, center, slot) -- the exact accumulation order of the old + # padded builder -- so the destination scatters reassociate identically + # and existing parity tolerances hold. The dense body is no longer a + # second copy of the edge math: ``call`` and ``call_graph`` share the + # one graph-native owner via ``_run_graph``. + nf, nloc, _ = nlist.shape + # Geometry enters in compute precision, as the old dense Step 1 did: + # edge vectors must be SUBTRACTED in compute precision (inside the + # converter), not computed in fp64 and cast after -- the two round + # differently for fp32 models. coord_ext = xp.astype(coord_ext, get_xp_precision(xp, self.compute_precision)) - nf, nloc, nnei = nlist.shape - nall = coord_ext.shape[1] - n_nodes = nf * nloc charge_spin = self._canonicalize_charge_spin( charge_spin, nf=nf, - dtype=coord_ext.dtype, - device=device, + ref=coord_ext, ) - - # === Step 2. Excluded type pairs === - if self.exclude_types: - # (nf, nloc, nnei), True means keep. - pair_keep_mask = xp.astype( - self.emask.build_type_exclude_mask(nlist, atype_ext), xp.bool - ) - else: - pair_keep_mask = xp.ones_like(nlist, dtype=xp.bool) - - # === Step 3. Type embedding (l=0) === - atype_loc = xp_take_first_n(atype_ext, 1, nloc) # (nf, nloc) - type_ebed = xp.reshape( - self.type_embedding(atype_loc), (n_nodes, self.channels) - ) # (N, C) - if self.charge_spin_embedding is not None: - type_ebed = self._apply_charge_spin_embedding( - type_ebed, - charge_spin, - nf=nf, - nloc=nloc, - ) - - # Native spin: condition the l=0 type features on the spin magnitude - # and hold the l=1 direction coefficients for the backbone seed. - spin_vec = None - if self.spin_embedding is not None and spin is not None: - type_ebed, spin_vec = self._apply_spin_embedding( - type_ebed, spin, xp.reshape(atype_loc, (-1,)), n_nodes=n_nodes - ) - - # === Step 4. Build edge cache once (geometry + RBF + Wigner-D) === - # Zone bridging (InnerClamp + SFPG + ZBL) is not routed through the - # standard DeePMD path: bridging only makes physical sense when - # paired with the ZBL energy that ``SeZMModel`` injects on the - # sparse-edge path, so ``forward`` keeps the original - # bridging-free aggregation semantics. - edge_cache = build_edge_cache( - type_ebed=type_ebed, - extended_coord=coord_ext, - nlist=nlist, - mapping=mapping, - pair_keep_mask=pair_keep_mask, - eps=self.eps, - deg_norm_floor=(self.deg_norm_floor if self.version >= 1.1 else self.eps), - edge_envelope=self.edge_envelope, - radial_basis=self.radial_basis, - n_radial=self.radial_basis.n_radial, - # Random local-Z roll is a training-only augmentation; - # the model is roll-equivariant, so inference fixes gamma. - random_gamma=False, - wigner_calc=self.wigner_calc, - build_wigner=self._need_full_wigner, + # The dense (nlist) lower has no comm implementation of its own and + # never will: it is the one owner of that rejection, so it raises + # here, before the graph build, instead of forwarding ``comm_dict`` + # into the shared trunk (which now threads comm through to the + # graph-native leaf for the graph route). + if comm_dict is not None: + raise NotImplementedError( + "the DPA4 dense (nlist) lower does not implement comm_dict " + "multi-rank exchange; freeze with the graph lower for " + "multi-rank inference" + ) + # ``_graph_from_padded_nlist`` owns the dense-topology contract: the + # pt-mirroring ``src_ok`` sanitization (out-of-range / broken-mapping + # sources masked) and, via ``graph_from_dense_quartet``, the + # ``mapping is None`` identity case (neighbor indices already local). + graph, atype_flat = _graph_from_padded_nlist( + coord_ext, atype_ext, nlist, mapping ) - - ebed_dim_0 = self.node_init_dim # (node_init_lmax+1)^2 - x0 = type_ebed # (N, C) - x0_out = x0 # (N, C) - - # === Step 5. Compute radial features once (fp32+) === - # Shape: (E, (node_init_lmax+1)*C) -> (E, node_init_lmax+1, C) - radial_feat = xp.reshape( - self.radial_embedding(edge_cache.edge_rbf), - (-1, self.node_init_lmax + 1, self.channels), - ) # (E, node_init_lmax+1, C) - if self.version >= 1.1: - radial_feat = radial_feat * xp.reshape(edge_cache.edge_env, (-1, 1, 1)) - - # === Step 6. Env FiLM conditioning (optional, fp32+) === - if self.use_env_seed: - atype_flat = xp.reshape(atype_loc, (-1,)) # (N,) - spin_flat = ( - xp.reshape(spin, (n_nodes, 3)) - if (self.spin_embedding is not None and spin is not None) - else None - ) - film = self.env_seed_embedding( - edge_cache=edge_cache, - atype_flat=atype_flat, - n_nodes=n_nodes, - spin=spin_flat, - ) # (N, 2*C) - scale_logits = film[:, : self.channels] # (N, C) - shift_logits = film[:, self.channels :] # (N, C) - scale_hat = ( - self.film_scale_norm(scale_logits) if self.edge_norm else scale_logits - ) # (N, C) - shift_hat = ( - self.film_shift_norm(shift_logits) if self.edge_norm else shift_logits - ) # (N, C) - scale_strength = xp.exp( - xp_asarray_nodetach( - xp, self.film_scale_strength_log[...], device=device - ) - ) - shift_strength = xp.exp( - xp_asarray_nodetach( - xp, self.film_shift_strength_log[...], device=device - ) - ) - scale = 1.0 + scale_strength * xp.tanh(scale_hat) # (N, C) - shift = shift_strength * xp.tanh(shift_hat) # (N, C) - x0_out = x0 * scale + shift - - # === Step 7. Build backbone l=0 features === - x = xp.concat( - [ - xp.reshape(x0_out, (n_nodes, 1, 1, self.channels)), - xp.zeros( - (n_nodes, ebed_dim_0 - 1, 1, self.channels), - dtype=type_ebed.dtype, - device=device, - ), - ], - axis=1, - ) # (N, D, 1, C) - - # === Step 8. Geometric Initial Embedding (+ neighbor spin l=1) === - if self.use_gie: - # GIE only needs l>=1, slice radial_feat[:, 1:, :] - zonal_coupling = self._build_gie_zonal_coupling(edge_cache) - spin_l1_message = ( - self.spin_embedding.edge_l1( - xp.reshape(spin, (n_nodes, 3)), - xp.reshape(atype_loc, (-1,)), - edge_cache, - ) - if (self.spin_embedding is not None and spin is not None) - else None - ) - x = ( - x - + self.gie( - n_nodes=n_nodes, - edge_cache=edge_cache, - radial_feat=radial_feat[:, 1:, :], - zonal_coupling=zonal_coupling, - spin_l1_message=spin_l1_message, - )[:, :, None, :] - ) - - # === Step 9. Add the on-site native spin l=1 to the backbone === - # The neighbor-spin l=1 is aggregated inside GIE (degree-normalized like - # the geometry); the atom's own spin direction is added here, un-normalized. - if spin_vec is not None: - spin_l1_rows = xp_asarray_nodetach(xp, self._spin_l1_rows, device=device) - spin_l1_src = spin_vec[:, :, None, :] # (N, 3, 1, C) - scatter_index = xp.broadcast_to( - xp.reshape(spin_l1_rows, (1, 3, 1, 1)), spin_l1_src.shape - ) - x = xp_scatter_sum(x, 1, scatter_index, spin_l1_src) - - # === Step 10. Fuse edge type features into radial features (fp32+) === - radial_feat = radial_feat + xp.reshape( - edge_cache.edge_type_feat, (-1, 1, self.channels) + x_scalar, _ = self._run_graph( + graph, + atype_flat, + nf=nf, + n_out_nodes=nf * nloc, + force_embedding=force_embedding, + charge_spin=charge_spin, + spin=spin, ) - radial_feat = xp.astype(radial_feat, get_xp_precision(xp, self.precision)) - rad_feat_per_block = [ - radial_feat[:, :rad_len, :] for rad_len in self.rad_sizes_per_block - ] # list of (E, lmax+1, C) - - # === Step 11. Convert to self.dtype and run blocks === - # The block stage is skipped entirely when there are no interaction - # blocks (zero-block descriptor) or no valid edges, sparing the working - # edge-cache dtype cast that only the blocks consume. - x = xp.astype(x, get_xp_precision(xp, self.precision)) # (N, D, 1, C) - if force_embedding is not None: - x = x + xp.astype(force_embedding, get_xp_precision(xp, self.precision)) - if self.blocks and edge_cache.src.shape[0] > 0: - edge_cache = edge_cache_to_dtype( - edge_cache, get_xp_precision(xp, self.precision) - ) - x = self._forward_blocks(x, edge_cache, rad_feat_per_block) - - # === Step 12. Final l=0 output mixing === - x_scalar = self._apply_readout(x, n_nodes) - - # === Step 13. Reshape to (nf, nloc, channels) and return === - descriptor = xp.reshape(x_scalar, (nf, nloc, self.channels)) # (nf, nloc, C) + # ``_run_graph`` returns (nf*nloc, 1, 1, channels) already in + # global precision; flatten the SO(3) singleton axes to (nf, nloc, C). + descriptor = xp.reshape(x_scalar, (nf, nloc, self.channels)) return ( - xp.astype(descriptor, get_xp_precision(xp, "global")), + descriptor, None, None, None, None, ) - def call_with_edges( + def call_graph( self, + graph: NeighborGraph, + atype: Array, + type_embedding: Array | None = None, + comm_dict: dict[str, Array] | None = None, + spin: Array | None = None, + charge_spin: Array | None = None, + ) -> tuple[Array, None]: + """Graph-native descriptor forward on the flat node axis. + + Parameters + ---------- + graph + Neighbor graph for the local atoms (ghost-free). ``edge_vec`` is + the geometry/autograd leaf; ``edge_mask`` flags valid edges. + atype + Flat node types with shape (N,). + type_embedding + Accepted for graph-seam interface stability and ignored: DPA4 + embeds types internally (``SeZMTypeEmbedding``) from ``atype``. + comm_dict + Border-exchange tensors for parallel inference, threaded down to + the interaction blocks. The dpmodel backend implements no + cross-rank exchange on any lower path: a block that actually + needs it raises from the ``exchange_ghost_features`` leaf (see + ``_run_graph``), not here. + spin + Per-node spin vectors with shape (N, 3) on the flat node axis, or + None. Consumed by ``spin_embedding`` (l=0 magnitude into the type + embedding, l=1 into the backbone and per-edge source features). + Ghost-free graphs need only per-local-atom spin. + charge_spin + Frame-level charge/spin conditioning with shape ``(nf, 2)`` (or a + shape ``_canonicalize_charge_spin`` can broadcast to it), or + ``None``. This is the SAME per-descriptor canonicalization the + dense ``call`` adapter applies (default-fill from + ``default_chg_spin`` when configured, shape validation, + broadcast to ``nf``); ``call_graph`` is the one owner of that + step on the graph route. ``nf`` is recovered from + ``graph.n_node.shape[0]`` (a static shape, safe under + ``torch.export``); each frame's node block must therefore hold + exactly ``N // nf`` nodes, which single-rank carry-all graphs + built from a rectangular ``(nf, nloc)`` input always satisfy. + + Returns + ------- + tuple[Array, None] + Flat ``(N, channels)`` descriptor in global precision, and + ``None`` (DPA4 produces no equivariant rot_mat for the fitting). + + Raises + ------ + NotImplementedError + When ``comm_dict`` is provided and a block needs the cross-rank + exchange (raised by the per-block leaf). + """ + n_nodes = atype.shape[0] + nf = graph.n_node.shape[0] + charge_spin = self._canonicalize_charge_spin( + charge_spin, + nf=nf, + ref=graph.edge_vec, + ) + x_scalar, _ = self._run_graph( + graph, atype, nf=nf, charge_spin=charge_spin, spin=spin, comm_dict=comm_dict + ) + # ``_run_graph`` returns the read-out with its SO(3) singleton + # axes still attached, shape (n_nodes, 1, 1, channels); flatten to the + # graph-seam contract shape (n_nodes, channels). + xp = array_api_compat.array_namespace(x_scalar) + x_scalar = xp.reshape(x_scalar, (n_nodes, self.channels)) + return x_scalar, None + + def _run_graph( + self, + graph: NeighborGraph, + atype_flat: Array, *, - coord_ext: Array, - atype_ext: Array, - edge_index: Array, - edge_vec: Array, - edge_mask: Array, + nf: int = 1, + n_out_nodes: int | None = None, force_embedding: Array | None = None, charge_spin: Array | None = None, spin: Array | None = None, comm_dict: dict[str, Array] | None = None, - nloc: int | None = None, ) -> tuple[Array, Array]: - """ - Compute the descriptor from a sparse edge list. - - Two node-set conventions share this path. In the single-domain path - (``comm_dict`` is ``None``) the nodes are exactly the local atoms and - ``edge_index`` source/destination both index ``[0, nf*nloc)``. In the - parallel (LAMMPS multi-rank) path the nodes span the extended region - (local owners followed by ghosts), ``edge_index`` indexes the extended - atoms directly, and each interaction block refreshes ghost-node features - from their owner ranks at the SO(2) convolution input (see - :func:`~deepmd.pt.model.descriptor.sezm_nn.block.exchange_ghost_features`). + """Graph-native descriptor forward shared by both descriptor entries. + + Both public entries -- the dense-nlist ``call`` adapter and the + graph-native ``call_graph`` -- funnel through here, so the descriptor's + graph-level pre-math work (its own ``exclude_types`` masking and + ``comm_dict`` threading) lives in exactly one place instead of being + duplicated per entry, then runs the edge-native core. + + The exclusion applied here is the DESCRIPTOR's own ``exclude_types`` + (via ``self.emask``), masked once onto the graph's ``edge_mask``; the + edge-cache core below has no exclusion parameter and never re-applies + it. This is a DIFFERENT knob from the MODEL-level ``pair_exclude_types``, + which is a graph-BUILD transform already folded into the incoming + graph/nlist (``make_model._call_common_graph`` / the NeighborList + builders / C++ ``applyPairExclusion``) and is never re-applied here. + + ``comm_dict`` is threaded unchanged to the interaction blocks; the + dpmodel backend implements no cross-rank exchange, so a block that + actually needs it raises from its per-block ``exchange_ghost_features`` + leaf (pt_expt overrides that leaf with a real ``border_op``). The dense + ``call`` adapter rejects ``comm_dict`` before it ever reaches here. Parameters ---------- - coord_ext - Coordinates with shape (nf, n*3) or (nf, n, 3) in Å, where ``n`` is - ``nloc`` in the single-domain path and ``nall`` in the parallel path. - atype_ext - Atom types with shape (nf, n). In the parallel path this spans the - extended region so ghost type embeddings are available for the - edge-type and environment-seed features. - edge_index - Edge indices with shape (2, E). - edge_vec - Edge vectors with shape (E, 3) in Å. - edge_mask - Edge mask with shape (E,). + graph + Neighbor graph for the local atoms; ``edge_vec`` is the + geometry/autograd leaf, ``edge_mask`` flags valid edges. + atype_flat + Flat node types with shape (N,). + nf + Frame count (only consumed by the charge/spin FiLM conditioning). + n_out_nodes + Leading node count kept for the read-out (owned atoms). ``None`` + keeps all nodes (``atype_flat.shape[0]``). force_embedding - Optional precomputed equivariant force embedding with shape - ``(nf * nloc, D, 1, channels)``, where - ``D = (node_init_lmax + 1) ** 2``. This tensor is added to the - initial SO(3) backbone state before the interaction blocks. + Optional per-node force conditioning, shape (N, D, 1, C). charge_spin - Frame-level charge and spin conditions with shape (nf, 2). + Optional charge/spin conditioning. + spin + Optional per-node spin vectors, shape (N, 3). comm_dict - Border-exchange tensors for parallel inference. When provided, the - node set spans the extended region and ghost features are exchanged - via ``deepmd_export::border_op`` between interaction blocks. - nloc - Number of owned (local) atoms per frame. Required when ``comm_dict`` - is provided; the final scalar read-out is restricted to these atoms. + MPI communication metadata forwarded to the interaction blocks; + ``None`` for single-rank inference. Returns ------- tuple[Array, Array] - The scalar descriptor with shape ``(nf, nloc, channels)`` and the - final equivariant latent with shape ``(nf * nloc, D_final, 1, channels)``. + Read-out with the SO(3) singleton axes still attached, shape + ``(n_out_nodes, 1, 1, channels)``, in global precision, and the + full multipole feature tensor ``x``. The public entries flatten + the singleton axes. + + Raises + ------ + NotImplementedError + When ``comm_dict`` is provided and a block actually needs the + cross-rank exchange (raised by the per-block leaf, not here). """ - xp = array_api_compat.array_namespace(coord_ext, atype_ext, edge_vec) - device = array_api_compat.device(coord_ext) - # === Step 1. Setup dimensions === - # ``n_per_frame`` is the per-frame node count: ``nloc`` in the - # single-domain path and ``nall`` in the parallel path. ``out_nloc`` is - # the owned-atom count used for the final local read-out. - coord_ext = xp.astype(coord_ext, get_xp_precision(xp, self.compute_precision)) - nf, n_per_frame = atype_ext.shape[:2] - parallel = comm_dict is not None - if parallel: - # Multi-rank parallel inference requires a custom border-exchange - # communication op that is not available in the dpmodel backend. - raise NotImplementedError( - "multi-rank comm_dict inference is not supported in the dpmodel backend" - ) - out_nloc = nloc if parallel else n_per_frame - atype_flat = xp.reshape(atype_ext, (-1,)) # (N,) + # Descriptor-owned exclusion: the descriptor's ``exclude_types`` masks + # the graph's ``edge_mask`` exactly once here; the edge-cache core has + # no exclusion parameter. (Model-level ``pair_exclude_types`` is a + # graph-BUILD transform already folded into the incoming graph.) + if self.exclude_types: + graph = apply_pair_exclusion(graph, atype_flat, self.emask) + if n_out_nodes is None: + n_out_nodes = atype_flat.shape[0] + edge_index = graph.edge_index + edge_vec = graph.edge_vec + edge_mask = graph.edge_mask + + xp = array_api_compat.array_namespace(edge_vec) + device = array_api_compat.device(edge_vec) # === Step 2. Type embedding (l=0) === type_ebed = xp.reshape( - self.type_embedding(atype_ext), (-1, self.channels) + self.type_embedding(atype_flat), (-1, self.channels) ) # (N, C) if self.charge_spin_embedding is not None: type_ebed = self._apply_charge_spin_embedding( type_ebed, charge_spin, nf=nf, - nloc=n_per_frame, + nloc=n_out_nodes // nf, ) n_nodes = type_ebed.shape[0] @@ -1549,9 +1542,8 @@ def call_with_edges( ) # === Step 3. Build edge cache once (sparse edges) === - edge_cache = build_edge_cache_from_edges( + edge_cache = _edge_cache_from_arrays( type_ebed=type_ebed, - atype_flat=atype_flat, edge_index=edge_index, edge_vec=edge_vec, edge_mask=edge_mask, @@ -1562,11 +1554,12 @@ def call_with_edges( bridging_switch=self.bridging_switch, edge_envelope=self.edge_envelope, radial_basis=self.radial_basis, - has_exclude_types=bool(self.exclude_types), - edge_type_keep_mask=self._edge_type_keep_mask, - # Random local-Z roll is a training-only augmentation; - # the model is roll-equivariant, so inference fixes gamma. - random_gamma=False, + # Random local-Z roll is a training-only augmentation; the model + # is roll-equivariant, so inference fixes gamma. Mirrors pt's + # ``random_gamma=self.random_gamma and self.training`` via the + # ``_in_training_mode`` runtime hook (False here; the pt_expt + # wrapper overrides it with the torch module's training flag). + random_gamma=self.random_gamma and self._in_training_mode(), wigner_calc=self.wigner_calc, build_wigner=self._need_full_wigner, ) @@ -1698,17 +1691,12 @@ def call_with_edges( # equals the whole node set and the slice is a no-op. Parallel # (single-frame): it drops the trailing ghost rows that only fed message # passing -- LAMMPS orders owned atoms before ghosts, so they lead. - n_out_nodes = nf * out_nloc - x = x[:n_out_nodes] + x = x[:n_out_nodes, ...] # === Step 12. Final l=0 output mixing === x_scalar = self._apply_readout(x, n_out_nodes) - # === Step 13. Reshape to (nf, nloc, channels) and return === - descriptor = xp.reshape( - x_scalar, (nf, out_nloc, self.channels) - ) # (nf, nloc, C) - return xp.astype(descriptor, get_xp_precision(xp, "global")), x + return xp.astype(x_scalar, get_xp_precision(xp, "global")), x def _forward_blocks( self, @@ -1879,7 +1867,7 @@ def _edge_quaternion(self, edge_cache: EdgeCache) -> Array: Parameters ---------- - edge_cache : EdgeFeatureCache + edge_cache : EdgeCache Per-edge cache. ``edge_quat`` is populated by the cache builder; the fallback covers caches produced without it. @@ -2009,42 +1997,6 @@ def _apply_spin_embedding( type_ebed = type_ebed + xp.astype(scalar, type_ebed.dtype) return type_ebed, vector - def _edge_type_keep_mask( - self, - atype_flat: Array, - src: Array, - dst: Array, - ) -> Array: - """ - Build keep mask for edge pairs based on excluded type pairs. - - Parameters - ---------- - atype_flat - Flattened local atom types with shape (N,). - src - Source indices with shape (E,). - dst - Destination indices with shape (E,). - - Returns - ------- - Array - Boolean mask with shape (E,), True means keep. - """ - xp = array_api_compat.array_namespace(atype_flat, src, dst) - if len(self.emask.exclude_types) == 0: - return xp.ones_like(src, dtype=xp.bool) - device = array_api_compat.device(atype_flat) - type_i = xp.take(atype_flat, dst, axis=0) - type_j = xp.take(atype_flat, src, axis=0) - type_i = xp.where(type_i >= 0, type_i, self.ntypes) - type_j = xp.where(type_j >= 0, type_j, self.ntypes) - type_ij = type_i * (self.ntypes + 1) + type_j - type_mask = xp_asarray_nodetach(xp, self.emask.type_mask[...], device=device) - keep = xp.take(type_mask, xp.astype(type_ij, xp.int64), axis=0) - return xp.astype(keep, xp.bool) - @staticmethod def _broadcast_grid_setting( value: bool | int | list[bool] | list[int], @@ -2192,8 +2144,7 @@ def _canonicalize_charge_spin( charge_spin: Array | None, *, nf: int, - dtype: Any, - device: Any, + ref: Array, ) -> Array | None: """ Canonicalize charge/spin conditions for the public descriptor path. @@ -2204,10 +2155,11 @@ def _canonicalize_charge_spin( Optional frame-level charge and spin conditions. nf Number of frames. - dtype - Target floating-point dtype. - device - Target device. + ref + A reference array from the caller's compute context; the + target namespace, dtype and device are all inferred from it + (array-API pitfall: deriving the namespace from the numpy + ``default_chg_spin`` attribute breaks the torch path). Returns ------- @@ -2216,17 +2168,19 @@ def _canonicalize_charge_spin( """ if self.charge_spin_embedding is None: return None + xp = array_api_compat.array_namespace(ref) + dtype = ref.dtype + device = array_api_compat.device(ref) if charge_spin is None: if self.default_chg_spin is None: raise ValueError("`charge_spin` is required for this SeZM descriptor.") - default_chg_spin = np.asarray(self.default_chg_spin) - xp = array_api_compat.array_namespace(default_chg_spin) charge_spin = xp.reshape( - xp_asarray_nodetach(xp, default_chg_spin, dtype=dtype, device=device), + xp_asarray_nodetach( + xp, np.asarray(self.default_chg_spin), dtype=dtype, device=device + ), (1, 2), ) else: - xp = array_api_compat.array_namespace(charge_spin) charge_spin = xp.astype(charge_spin, dtype) if charge_spin.ndim == 1: @@ -2324,17 +2278,81 @@ def has_message_passing(self) -> bool: return True def has_message_passing_across_ranks(self) -> bool: - """Whether multi-rank inference needs cross-rank ghost-feature exchange. - - SeZM reads ghost-neighbour features at every interaction block, so a - domain-decomposed run must exchange them through ``border_op``. Source - Freeze Propagation bridging is excluded: its per-node gate folds a - node's entire outgoing-edge set, which a single rank cannot observe for - ghost owners, so the edge-based with-comm artifact is not exported for - bridging models and multi-rank inference fails fast instead. + """Whether multi-rank inference needs cross-rank ghost exchange. + + SeZM reads ghost-neighbour features at every interaction block; the + GRAPH lower implements the exchange via per-block ``border_op`` + (pt_expt ``exchange_ghost_features``). Source Freeze Propagation + bridging is excluded: its per-node gate folds a node's entire + outgoing-edge set, which a single rank cannot observe for ghost + owners, so bridging models fail fast on multi-rank instead. + + The DENSE (nlist) lower remains comm-less — see + :meth:`dense_lower_supports_comm`; the freeze machinery consults both + so nlist-kind artifacts carry ``has_comm_artifact=False``. """ return self.bridging_switch is None + def dense_lower_supports_comm(self) -> bool: + """The DPA4 dense (nlist) lower has no comm_dict implementation. + + The dense adapter raises on ``comm_dict``; only the graph lower + exchanges ghosts. Consulted by the freeze machinery for non-graph + lower kinds so no dead-comm dense artifact is ever emitted. + """ + return False + + def uses_graph_lower(self) -> bool: + """Whether this descriptor supports the sel-free NeighborGraph lower. + + Returns + ------- + bool + False only when the escape hatch has been pulled + (``disable_graph_lower()`` / ``_graph_lower_disabled``). Every + conditioning input DPA4 supports -- native spin + (``spin_embedding``), charge/spin FiLM (``charge_spin_embedding``), + and SFPG bridging (``bridging_switch``) -- rides the graph lower: + spin and charge_spin are threaded through ``call_graph`` like any + other per-node/per-frame input, and bridging is applied inside + the shared ``_run_graph`` forward with no extra threading (it + reads ``self.bridging_switch`` directly). Bridging models still + fail multi-rank fast via ``has_message_passing_across_ranks``. + """ + return not self._graph_lower_disabled + + def uses_compact_edge_pairs(self) -> bool: + """DPA4 attention is a per-edge scatter softmax; no pair axis.""" + return False + + def _in_training_mode(self) -> bool: + """Whether the descriptor is currently in training mode. + + Gates the training-only random local-Z roll + (``random_gamma``): the roll is applied only when the descriptor is + training, mirroring pt's ``self.random_gamma and self.training``. + dpmodel is an inference/reference backend with no training mode, so + this returns ``False`` (deterministic, fixed gamma); the pt_expt + wrapper overrides it with the torch module's ``training`` flag. + """ + return False + + def supports_native_spin(self) -> bool: + """DPA4 accepts a per-node ``spin`` on ``call_graph`` (native magnetic conditioning); overrides the ``BaseDescriptor`` default of ``False``.""" + return True + + def supports_charge_spin(self) -> bool: + """DPA4 accepts a frame-level ``charge_spin`` on ``call_graph`` (FiLM conditioning); overrides the ``BaseDescriptor`` default of ``False``.""" + return True + + def disable_graph_lower(self) -> None: + """Route this descriptor through the legacy dense lower.""" + self._graph_lower_disabled = True + + def graph_type_embedding_table(self) -> None: + """DPA4 embeds types internally; the graph seam passes nothing.""" + return None + def need_sorted_nlist_for_lower(self) -> bool: return False diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/__init__.py b/deepmd/dpmodel/descriptor/dpa4_nn/__init__.py index 54d1e3045c..85c854e69c 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/__init__.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/__init__.py @@ -30,8 +30,6 @@ ) from .edge_cache import ( EdgeCache, - build_edge_cache, - build_edge_cache_from_edges, build_edge_type_feat, compute_edge_src_gate, edge_cache_to_dtype, @@ -164,8 +162,6 @@ "WignerDCalculator", "apply_lora_to_sezm", "build_cartesian_basis", - "build_edge_cache", - "build_edge_cache_from_edges", "build_edge_cartesian_tensors", "build_edge_quaternion", "build_edge_type_feat", diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/attention.py b/deepmd/dpmodel/descriptor/dpa4_nn/attention.py index fd2829b6e3..a06e9fe2be 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/attention.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/attention.py @@ -63,8 +63,8 @@ def segment_envelope_gated_softmax( denominator sum are scattered over these indices, which makes the normalization layout-agnostic: it is correct both for the padded ``call`` (where ``dst == repeat(arange(n_nodes), nnei)``) and for the - sparse ``call_with_edges`` (arbitrary ``dst`` order and per-node - degree). + graph-native ``call_graph`` route (arbitrary ``dst`` order and + per-node degree). n_nodes Number of nodes. z_bias_raw @@ -150,6 +150,9 @@ def segment_envelope_gated_softmax( # === Step 2. Destination-wise max including the physical null mass === # The null initialization keeps empty and all-masked segments finite. + # Destination segment max over ``dst`` (pt ``scatter_reduce`` amax) is + # layout-agnostic and order-independent, so the padded ``call`` stays + # bit-exact while the graph-native ``call_graph`` route shares this path. group_max = xp_maximum_at( xp.zeros((n_nodes, n_channel), dtype=compute_dtype, device=device) + null_logit, dst, diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/block.py b/deepmd/dpmodel/descriptor/dpa4_nn/block.py index 751efe3f95..3d12a3be3e 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/block.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/block.py @@ -66,7 +66,7 @@ ) from .edge_cache import ( - EdgeFeatureCache, + EdgeCache, ) @@ -644,7 +644,7 @@ def __init__( def call( self, x: Array, - edge_cache: EdgeFeatureCache, + edge_cache: EdgeCache, radial_feat: Array, unit_history: list[Array] | None = None, comm_dict: dict[str, Array] | None = None, @@ -709,7 +709,7 @@ def _extract_l0_from_canonical(self, value: Array) -> Array: def _run_so2_unit( self, x: Array, - edge_cache: EdgeFeatureCache, + edge_cache: EdgeCache, radial_feat: Array, comm_dict: dict[str, Array] | None = None, ) -> Array: @@ -742,7 +742,7 @@ def _run_so2_unit( def _run_so2_unit_impl( self, x: Array, - edge_cache: EdgeFeatureCache, + edge_cache: EdgeCache, radial_feat: Array, ) -> Array: """Run the SO(2) unit implementation.""" @@ -804,7 +804,7 @@ def _run_ffn_unit_impl(self, x: Array, unit_idx: int) -> Array: def _forward_with_residual_shortcuts( self, x: Array, - edge_cache: EdgeFeatureCache, + edge_cache: EdgeCache, radial_feat: Array, unit_history: list[Array] | None = None, comm_dict: dict[str, Array] | None = None, @@ -851,7 +851,7 @@ def _forward_with_residual_shortcuts( def _forward_with_full_attn_res( self, x: Array, - edge_cache: EdgeFeatureCache, + edge_cache: EdgeCache, radial_feat: Array, unit_history: list[Array] | None = None, comm_dict: dict[str, Array] | None = None, @@ -915,7 +915,7 @@ def _forward_with_full_attn_res( def _forward_with_block_attn_res( self, x: Array, - edge_cache: EdgeFeatureCache, + edge_cache: EdgeCache, radial_feat: Array, unit_history: list[Array] | None = None, comm_dict: dict[str, Array] | None = None, diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py b/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py index 3ba5dde121..99f20c4285 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py @@ -27,11 +27,11 @@ ) import array_api_compat -import numpy as np from deepmd.dpmodel.array_api import ( xp_add_at, xp_asarray_nodetach, + xp_uniform, ) from .utils import ( @@ -44,7 +44,6 @@ ) WignerCalculatorFn = Callable[[Any], "tuple[Any, Any]"] -EdgeTypeKeepMaskFn = Callable[[Any, Any, Any], Any] @dataclass @@ -104,7 +103,7 @@ class EdgeCache: Validity mask for the padded standard-path layout with shape (E,) or (E, 1); 1 marks a real edge, 0 a padded/invalid slot. ``None`` means all slots are valid (e.g. the sparse - :func:`build_edge_cache_from_edges` path, where masking is folded into + :func:`_edge_cache_from_arrays` path, where masking is folded into the per-edge weights). This field has no pt counterpart. """ @@ -228,204 +227,9 @@ def compute_edge_src_gate( return xp.take(eta, src, axis=0)[:, None] -def build_edge_cache( +def _edge_cache_from_arrays( *, type_ebed: Any, - extended_coord: Any, - nlist: Any, - mapping: Any, - pair_keep_mask: Any, - eps: float, - deg_norm_floor: float, - edge_envelope: Callable[[Any], Any], - radial_basis: Callable[[Any], Any], - n_radial: int, # unused in padded layout; kept for pt signature parity - random_gamma: bool, - wigner_calc: WignerCalculatorFn, - build_wigner: bool = True, - gamma: Any = None, -) -> EdgeCache: - """ - Build the global edge cache from a DeePMD padded neighbor list. - - This converts DeePMD's per-frame padded neighbor list into the per-edge - tensors reused across blocks. Where pt extracts a sparse list of valid - edges with ``torch.nonzero`` (data-dependent length), the array-API port - keeps one edge slot for every neighbor slot, so ``E = nf * nloc * nnei`` - flattened row-major and ``dst == repeat(arange(nf * nloc), nnei)``. Invalid - slots (``nlist == -1`` padding, excluded type pairs, out-of-range mapped - sources) stay in the arrays, flagged by ``edge_mask``; their geometry, - envelope, radial basis, and type features are masked to zero. - - The resulting cache contains: - - - per-edge endpoints: ``src``, ``dst`` and per-edge type features: ``edge_type_feat`` (src+dst) - - per-edge geometry: ``edge_vec`` - - per-edge smooth weights: C^3 cutoff envelope ``edge_env`` - - per-edge radial basis: ``edge_rbf`` (envelope already baked in) - - per-edge rotation blocks: block-diagonal Wigner-D matrices ``D_full`` and ``Dt_full`` - - destination-node smooth normalization: ``inv_sqrt_deg`` from - envelope-squared degree ``sum(edge_env**2)`` - - Notes - ----- - Input formats follow DeePMD conventions: - - - ``extended_coord`` has shape ``(nf, nall, 3)``. - - ``nlist`` has shape ``(nf, nloc, nnei)`` and stores indices into the extended axis - (``0..nall-1``), with ``-1`` indicating padding. - - ``mapping`` (when provided) maps extended indices to local indices ``0..nloc-1``. - When ``mapping`` is ``None``, the function assumes the neighbor indices are already local. - - Gathered edge vectors on invalid slots are garbage (placeholder index 0) - and may even be exactly zero (self-difference), which would produce a 0/0 - in the normalization inside the quaternion construction. Although the - forward contribution of such slots is masked out downstream, a NaN there - would still poison the backward pass (``where`` propagates NaN gradients - from the unselected branch). Invalid slots are therefore rewritten to the - safe dummy unit vector ``+z`` before any norm/quaternion/Wigner evaluation, - and their envelope, radial basis, and type features are multiplied by the - mask so they are exactly zero. - - Parameters - ---------- - type_ebed - Per-node type embedding with shape (N, C), where N=nf*nloc. - extended_coord - Extended coordinates with shape (nf, nall, 3). - nlist - Neighbor list with shape (nf, nloc, nnei). - mapping - Mapping from extended indices to local indices with shape (nf, nall), or None. - pair_keep_mask - Pair keep mask from `PairExcludeMask` with shape (nf, nloc, nnei). True means keep. - eps - Small positive epsilon for safe norm. - deg_norm_floor - Floor added to the envelope-squared degree before inverse-sqrt - normalization. - edge_envelope - C^3 edge envelope module. - radial_basis - Radial basis module. - n_radial - Number of radial basis channels. Unused here; kept for signature - parity with pt. - random_gamma - Whether to apply a random roll around the local +Z axis before - constructing Wigner-D blocks. - wigner_calc - Callable that converts edge-aligned quaternions into packed Wigner-D - blocks. - gamma - Optional per-edge roll angles with shape (E,), used only when - ``random_gamma`` is True. pt draws gamma internally with - ``torch.rand`` and the draw cannot be reproduced here, so callers - needing determinism (e.g. tests) inject the angles explicitly. When - None, angles are drawn from ``numpy.random.default_rng()`` uniformly - in ``[0, 2*pi)``, matching pt's distribution. - - Returns - ------- - EdgeCache - Padded per-edge cache with ``edge_mask`` set. - """ - xp = array_api_compat.array_namespace(type_ebed, extended_coord, nlist) - device = array_api_compat.device(extended_coord) - nf, nloc, nnei = nlist.shape - nall = extended_coord.shape[1] - n_nodes = nf * nloc - - # === Step 1. Validity mask and safe indices (pt edge_keep semantics) === - mask, nlist_safe, src_local_safe = _build_edge_mask_and_src( - xp, nlist, mapping, pair_keep_mask, nall - ) - mask_flat = xp.reshape(mask, (-1,)) - - # === Step 2. Node indices === - # dst is slot-implicit: arange(nf * nloc) repeated nnei times (contract). - frame_idx = xp.reshape(xp.arange(nf, dtype=xp.int64, device=device), (nf, 1, 1)) - src = xp.reshape(frame_idx * nloc + src_local_safe, (-1,)) - node_idx = xp.arange(n_nodes, dtype=xp.int64, device=device) - dst = xp.reshape(xp.broadcast_to(node_idx[:, None], (n_nodes, nnei)), (-1,)) - - # === Step 3. Gather per-edge geometry from extended coordinates === - # edge_vec points from center -> neighbor: r_ij = r_j - r_i (in Å). - coord_flat = xp.reshape(extended_coord, (nf * nall, 3)) - neighbor_coord_index = xp.reshape(frame_idx * nall + nlist_safe, (-1,)) - loc_idx = xp.reshape(xp.arange(nloc, dtype=xp.int64, device=device), (1, nloc, 1)) - center_ext = xp.broadcast_to(frame_idx * nall + loc_idx, (nf, nloc, nnei)) - center_coord_index = xp.reshape(center_ext, (-1,)) - neighbor_pos = xp.take(coord_flat, neighbor_coord_index, axis=0) - center_pos = xp.take(coord_flat, center_coord_index, axis=0) - vec = neighbor_pos - center_pos # (E, 3) - - # === Step 4. Rewrite invalid slots to the safe +z dummy vector === - # Gradient safety: see the function docstring. edge_len is the scalar - # distance, computed only after the safe rewrite so it stays finite. - maskf = xp.astype(mask_flat, vec.dtype)[:, None] # (E, 1) - z_unit = xp_asarray_nodetach( - xp, np.array([[0.0, 0.0, 1.0]]), dtype=vec.dtype, device=device - ) - edge_vec = vec * maskf + (1.0 - maskf) * z_unit - edge_len = safe_norm(edge_vec, eps) # (E, 1) - - # === Step 5. Envelope and radial basis, masked to zero on invalid slots === - # Edges with r >= rcut are not removed from the cache. Their envelope is - # exactly zero, so messages vanish naturally while degree normalization - # remains smooth at the cutoff boundary. - edge_env = edge_envelope(edge_len) * maskf # (E, 1) - edge_rbf = radial_basis(edge_len) * maskf # (E, n_radial) - - # === Step 6. Edge quaternion -> Wigner-D blocks === - D_full, Dt_full, edge_quat = _build_edge_wigner( - edge_vec=edge_vec, - edge_len=edge_len, - eps=eps, - random_gamma=random_gamma, - wigner_calc=wigner_calc, - gamma=gamma, - build_full=build_wigner, - ) # (E, D, D), (E, D, D), (E, 4) - - # === Step 7. Edge type features (src + dst), masked === - edge_type_feat = build_edge_type_feat(type_ebed, src, dst) * xp.astype( - maskf, type_ebed.dtype - ) # (E, C) - - # === Step 8. Smooth destination degrees === - # pt accumulates env^2 with ``index_add_`` over dst; in the padded - # node-contiguous layout this is a plain masked sum over the nnei axis. - # edge_env is already exactly zero on invalid slots. - env_sq = xp.reshape(edge_env[:, 0] * edge_env[:, 0], (n_nodes, nnei)) - deg = xp.sum(env_sq, axis=1) # (N,) - inv_sqrt_deg = xp.reshape( - 1.0 / xp.sqrt(deg + deg_norm_floor), (n_nodes, 1, 1) - ) # (N, 1, 1) - - return EdgeCache( - src=src, - dst=dst, - edge_type_feat=edge_type_feat, - edge_vec=edge_vec, - edge_rbf=edge_rbf, - edge_env=edge_env, - deg=deg, - inv_sqrt_deg=inv_sqrt_deg, - D_full=D_full, - Dt_full=Dt_full, - D_to_m_cache={}, - Dt_from_m_cache={}, - edge_src_gate=None, - edge_quat=edge_quat, - edge_mask=mask_flat, - ) - - -def build_edge_cache_from_edges( - *, - type_ebed: Any, - atype_flat: Any, edge_index: Any, edge_vec: Any, edge_mask: Any, @@ -436,8 +240,6 @@ def build_edge_cache_from_edges( bridging_switch: Callable[[Any], Any] | None, edge_envelope: Callable[[Any], Any], radial_basis: Callable[[Any], Any], - has_exclude_types: bool, - edge_type_keep_mask: EdgeTypeKeepMaskFn, random_gamma: bool, wigner_calc: WignerCalculatorFn, build_wigner: bool = True, @@ -446,12 +248,16 @@ def build_edge_cache_from_edges( """ Build the global edge cache from a sparse edge list. + Private core, invoked only from ``DescrptDPA4._run_graph``. The descriptor's + own ``exclude_types`` masking is not applied here: ``_run_graph`` applies it + exactly once, upstream, on the ``NeighborGraph``'s ``edge_mask`` via + ``apply_pair_exclusion``. (Model-level ``pair_exclude_types`` is a separate, + graph-BUILD transform, already folded into the incoming graph.) + Parameters ---------- type_ebed Per-node type embedding with shape (N, C), where N=nf*nloc. - atype_flat - Flattened local atom types with shape (N,). edge_index Edge indices with shape (2, E). edge_vec @@ -478,10 +284,6 @@ def build_edge_cache_from_edges( C^3 edge envelope module. radial_basis Radial basis module. - has_exclude_types - Whether excluded type pairs should be filtered in this path. - edge_type_keep_mask - Callable that builds the keep mask for edge type exclusions. random_gamma Whether to apply a random roll around the local +Z axis before constructing Wigner-D blocks. @@ -490,11 +292,9 @@ def build_edge_cache_from_edges( blocks. gamma Optional per-edge roll angles with shape (E,), used only when - ``random_gamma`` is True. pt draws gamma internally with - ``torch.rand`` and the draw cannot be reproduced here, so callers - needing determinism (e.g. tests) inject the angles explicitly. When - None, angles are drawn from ``numpy.random.default_rng()`` uniformly - in ``[0, 2*pi)``, matching pt's distribution. + ``random_gamma`` is True. When None, drawn with the backend's RNG + (:func:`~deepmd.dpmodel.array_api.xp_uniform`) uniformly in + ``[0, 2*pi)``; callers may inject angles to pin a draw. Returns ------- @@ -504,13 +304,11 @@ def build_edge_cache_from_edges( xp = array_api_compat.array_namespace(type_ebed, edge_index, edge_vec) device = array_api_compat.device(edge_vec) n_nodes = type_ebed.shape[0] - src = xp.astype(edge_index[0], xp.int64) - dst = xp.astype(edge_index[1], xp.int64) + src = xp.astype(edge_index[0, ...], xp.int64) + dst = xp.astype(edge_index[1, ...], xp.int64) - # === Step 1. Normalize mask and apply type exclusions === + # === Step 1. Normalize mask === edge_keep = xp.astype(edge_mask, xp.bool) - if has_exclude_types: - edge_keep = edge_keep & edge_type_keep_mask(atype_flat, src, dst) # === Step 2. Promote geometry dtype === edge_vec = xp.astype(edge_vec, compute_dtype) @@ -606,9 +404,9 @@ def _build_edge_wigner( blocks. gamma Optional per-edge roll angles with shape (E,), used only when - ``random_gamma`` is True. When None, angles are drawn from - ``numpy.random.default_rng()`` uniformly in ``[0, 2*pi)``, matching - pt's ``torch.rand`` distribution. + ``random_gamma`` is True. When None, drawn with the backend's RNG + (:func:`~deepmd.dpmodel.array_api.xp_uniform`) uniformly in + ``[0, 2*pi)``. build_full Whether to materialize the full ``(E, D, D)`` Wigner-D blocks. When False (all message-passing blocks take the Cartesian path), only the @@ -632,13 +430,11 @@ def _build_edge_wigner( ) # === Step 2. Apply optional random local-Z roll === - # pt draws the roll with ``torch.rand``; here it is injected or drawn from - # numpy so the array-API call site stays reproducible. + # Training-only augmentation: drawn with the backend's own RNG so torch + # replays it under setup_seed and keeps the draw on-device. if random_gamma: if gamma is None: - gamma = np.random.default_rng().uniform( - 0.0, 2.0 * math.pi, edge_quat.shape[0] - ) + gamma = xp_uniform(edge_quat, edge_quat.shape[0], 0.0, 2.0 * math.pi) gamma = xp.astype( xp_asarray_nodetach(xp, gamma, device=device), edge_quat.dtype ) @@ -736,77 +532,6 @@ def _finalize_edge_cache( ) -def _build_edge_mask_and_src( - xp: Any, - nlist: Any, - mapping: Any, - pair_keep_mask: Any, - nall: int, -) -> tuple[Any, Any, Any]: - """ - Build the padded edge validity mask and safe source-local indices. - - This reproduces the pt edge-keep rules for the padded layout: - - - padding slots (``nlist == -1``) are invalid; - - excluded type pairs (``pair_keep_mask == False``) are invalid; - - after mapping the neighbor's extended index to a local index, slots - whose source falls outside ``[0, nloc)`` are invalid (pt's ``src_ok`` - filter; e.g. broken mapping or ghost-only neighbors); - - no distance-based filtering: edges beyond ``rcut`` stay valid and are - zeroed naturally by the smooth envelope. - - Instead of dropping invalid slots (pt's ``torch.nonzero``), they are kept - with ``mask == False`` and safe (index 0) placeholder indices. - - Parameters - ---------- - xp - Array namespace. - nlist - Neighbor list with shape (nf, nloc, nnei); -1 marks padding. - mapping - Extended-to-local mapping with shape (nf, nall), or None if the - neighbor indices are already local. - pair_keep_mask - Pair exclusion keep mask with shape (nf, nloc, nnei). True means keep. - nall - Number of atoms on the extended axis per frame. - - Returns - ------- - tuple[Array, Array, Array] - ``(mask, nlist_safe, src_local_safe)``, all with shape - (nf, nloc, nnei). ``mask`` is boolean; the two index arrays are int64 - with 0 substituted on invalid slots. - """ - nf, nloc, nnei = nlist.shape - nlist = xp.astype(nlist, xp.int64) - mask = (nlist >= 0) & pair_keep_mask - nlist_safe = xp.where(mask, nlist, xp.zeros_like(nlist)) - - if mapping is None: - # Neighbor indices are already local indices in [0, nloc). - src_local = nlist_safe - else: - # Map extended index -> local index for each frame. - mapping_flat = xp.astype(xp.reshape(mapping, (-1,)), xp.int64) - frame_idx = xp.reshape( - xp.arange(nf, dtype=xp.int64, device=array_api_compat.device(nlist)), - (nf, 1, 1), - ) - flat_idx = xp.reshape(frame_idx * nall + nlist_safe, (-1,)) - src_local = xp.reshape(xp.take(mapping_flat, flat_idx, axis=0), nlist.shape) - - # pt's src_ok filter: drop (here: mask) edges mapping outside [0, nloc). - mask = mask & (src_local >= 0) & (src_local < nloc) - src_local_safe = xp.where(mask, src_local, xp.zeros_like(src_local)) - # Re-zero nlist_safe after the src_ok update so coordinate gathers stay - # in-bounds when callers pass local nlists with out-of-range entries. - nlist_safe = xp.where(mask, nlist_safe, xp.zeros_like(nlist_safe)) - return mask, nlist_safe, src_local_safe - - def build_edge_type_feat( type_ebed: Any, src: Any, diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py index 238f904a8c..b624126b40 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py @@ -330,8 +330,8 @@ def call( # applied after the validity masking below. This reduction is # layout-agnostic: it is correct both for the padded ``call`` (row-major # ``dst`` makes the accumulation order identical to a sum over the - # ``nnei`` axis, hence bit-exact) and for the sparse ``call_with_edges`` - # (arbitrary ``dst`` order and per-node degree). The l=0 row is left at + # ``nnei`` axis, hence bit-exact) and for the graph-native ``call_graph`` + # route (arbitrary ``dst`` order and per-node degree). The l=0 row is left at # its zero initialization by concatenating it below the contiguous # non-scalar rows 1..D-1. edge_mask = edge_cache.edge_mask @@ -686,8 +686,8 @@ def call( # Destination scatter-add over ``dst`` (pt ``index_add_``), applied after # the validity masking below. Layout-agnostic: correct for the padded # ``call`` (row-major ``dst`` keeps the accumulation order identical to a - # sum over the ``nnei`` axis, hence bit-exact) and for the sparse - # ``call_with_edges`` (arbitrary ``dst`` order and per-node degree). + # sum over the ``nnei`` axis, hence bit-exact) and for the graph-native + # ``call_graph`` route (arbitrary ``dst`` order and per-node degree). edge_mask = edge_cache.edge_mask if edge_mask is not None: outer_flat = outer_flat * xp.astype( diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index c1938a3b01..6594e4c632 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -137,6 +137,91 @@ def has_message_passing_across_ranks(self) -> bool: """ return False + def supports_native_spin(self) -> bool: + """Returns whether the descriptor natively conditions on per-atom spin. + + Declaring ``True`` obliges the descriptor's ``call_graph`` to + accept a per-node ``spin`` keyword; the atomic model only + forwards the keyword to descriptors that declare the capability, + since an unconditional ``spin=`` kwarg would be a ``TypeError`` + on a ``call_graph`` signature that does not declare it. + + Concrete default ``False`` so descriptors across all backends + (pt/pd/tf subclass this same base) need no change until they grow + a native spin mechanism of their own; such descriptors override + this method to return ``True``. + """ + return False + + def supports_charge_spin(self) -> bool: + """Returns whether the descriptor conditions on a frame-level ``charge_spin`` input. + + Declaring ``True`` obliges the descriptor's ``call_graph`` to + accept a frame-level ``charge_spin`` keyword; the atomic model + only forwards the keyword to descriptors that declare the + capability. Concrete default ``False`` (see + ``supports_native_spin``); descriptors that condition on this + input override this method to return ``True``. + """ + return False + + def uses_graph_lower(self) -> bool: + """Returns whether the descriptor supports the graph-native (NeighborGraph) lower. + + Declaring ``True`` obliges the descriptor to implement + ``call_graph``; the model layer routes ``forward_lower`` through + the NeighborGraph path only for descriptors that declare the + capability, and falls back to the legacy dense (nlist) lower + otherwise. + + Concrete default ``False`` so descriptors across all backends + (which subclass this same base) stay on the dense lower until + they implement a graph-native forward; such descriptors override + this method (typically conditioning on their configuration and + on :meth:`disable_graph_lower`). + """ + return False + + def disable_graph_lower(self) -> None: + """Force the legacy dense (nlist) lower for this descriptor. + + An explicit opt-out knob used by contexts where the graph-native + lower is unsupported or undesirable. After calling this, + :meth:`uses_graph_lower` must return ``False`` regardless of the + descriptor configuration. + + Concrete default: a no-op, since a descriptor without a graph + lower is already dense-only. Descriptors overriding + :meth:`uses_graph_lower` must also override this to set their + escape hatch. + """ + return None + + def uses_compact_edge_pairs(self) -> bool: + """Returns whether the descriptor's graph lower traces compact edge pairs. + + The compact ``center_edge_pairs`` realization uses + unbacked-SymInt ``nonzero``/``repeat`` sizes when traced for + export; ``check_graph_trace_torch_version`` keys its + torch >= 2.6 requirement on this capability. Concrete default + ``False``; only meaningful for descriptors whose + :meth:`uses_graph_lower` can return ``True``. + """ + return False + + def graph_type_embedding_table(self) -> Any | None: + """Full type-embedding table consumed by the graph-route forward. + + Returns + ------- + Any | None + The ``(ntypes + 1, tebd_dim)`` type-embedding table for + descriptors whose graph lower consumes an external table, or + ``None`` (the concrete default) for descriptors that embed + types internally or have no graph lower. + """ + return None + @abstractmethod def need_sorted_nlist_for_lower(self) -> bool: """Returns whether the descriptor needs sorted nlist when using `forward_lower`.""" diff --git a/deepmd/dpmodel/loss/ener_spin.py b/deepmd/dpmodel/loss/ener_spin.py index 4446172634..30959fa83f 100644 --- a/deepmd/dpmodel/loss/ener_spin.py +++ b/deepmd/dpmodel/loss/ener_spin.py @@ -232,8 +232,13 @@ def call( if self.has_fr: find_force = label_dict.get("find_force", 0.0) pref_fr = pref_fr * find_force - force_pred = model_dict["force"] - force_label = label_dict["force"] + # Reshape to the canonical (nf, natoms, 3) atomic shape: the raw + # data-loader label is flat (nf, natoms * 3), matching the + # ``xp.reshape(label_dict[...], (-1, natoms, ncomp))`` idiom used + # by every other atomic-label loss (see ``dpmodel/loss/dos.py`` + # and ``dpmodel/loss/tensor.py``). + force_pred = xp.reshape(model_dict["force"], (-1, natoms, 3)) + force_label = xp.reshape(label_dict["force"], (-1, natoms, 3)) if self.loss_func == "mse": diff_fr = force_label - force_pred # [nf, nloc, 3] if maskf is not None: @@ -276,8 +281,9 @@ def call( if self.has_fm: find_force_mag = label_dict.get("find_force_mag", 0.0) pref_fm = pref_fm * find_force_mag - force_mag_pred = model_dict["force_mag"] - force_mag_label = label_dict["force_mag"] + # Same flat -> (nf, natoms, 3) reshape as the real-force branch above. + force_mag_pred = xp.reshape(model_dict["force_mag"], (-1, natoms, 3)) + force_mag_label = xp.reshape(label_dict["force_mag"], (-1, natoms, 3)) mask_mag = model_dict["mask_mag"] # mask_mag: [nframes, natoms, 1], bool -> use mask multiplication mask_float = xp.astype(mask_mag, force_mag_pred.dtype) diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index 9da6f8e585..9c85fa7e26 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -98,6 +98,15 @@ def is_aparam_nall(self) -> bool: def model_output_type(self) -> list[str]: """Get the output type for the model.""" + def has_spin(self) -> bool: + """Returns whether the model has spin input and output. + + Concrete default ``False`` so non-spin models across all backends + (which subclass this same base) need no change; spin-capable + model classes override this method to return ``True``. + """ + return False + @abstractmethod def serialize(self) -> dict: """Serialize the model. diff --git a/deepmd/dpmodel/model/dp_linear_model.py b/deepmd/dpmodel/model/dp_linear_model.py new file mode 100644 index 0000000000..132dc48385 --- /dev/null +++ b/deepmd/dpmodel/model/dp_linear_model.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""dpmodel linear energy model: a make_model CM over the linear atomic-model +composition (twin of ``deepmd.pt_expt.model.dp_linear_model``). +""" + +from typing import ( + Any, +) + +from deepmd.dpmodel.atomic_model.linear_atomic_model import ( + LinearEnergyAtomicModel, +) +from deepmd.dpmodel.common import ( + NativeOP, +) +from deepmd.dpmodel.model.base_model import ( + BaseModel, +) +from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, +) +from deepmd.dpmodel.model.make_model import ( + make_model, +) + +DPLinearModel_ = make_model(LinearEnergyAtomicModel, T_Bases=(NativeOP, BaseModel)) + + +@BaseModel.register("linear_ener") # config type +@BaseModel.register("linear") # wire type emitted by the flat serialize +class LinearEnergyModel(DPModelCommon, DPLinearModel_): + r"""Energy model over a linear combination of atomic models. + + The atomic energy is the weighted sum of the children's atomic + energies; on the NeighborGraph route every child consumes the same + graph, so the summed energy differentiates through one shared edge + backward. Used e.g. for analytical bridging compositions + (learned model + :class:`~deepmd.dpmodel.atomic_model.inter_potential.InterPotentialAtomicModel`). + """ + + def __init__( + self, + *args: Any, + **kwargs: Any, + ) -> None: + DPModelCommon.__init__(self) + DPLinearModel_.__init__(self, *args, **kwargs) diff --git a/deepmd/dpmodel/model/dpa4_model.py b/deepmd/dpmodel/model/dpa4_model.py new file mode 100644 index 0000000000..f5978b7fd4 --- /dev/null +++ b/deepmd/dpmodel/model/dpa4_model.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""DPA4/SeZM-family energy model. + +A THIN subclass of the descriptor-agnostic :class:`EnergyModel` that owns +exactly the dpa4-family concerns: the family's registry wire types +(``dpa4_ener``/``sezm_ener`` fitting-type dispatch keys). The generic +``EnergyModel`` stays free of descriptor-specific registrations. +""" + +from deepmd.dpmodel.model.base_model import ( + BaseModel, +) +from deepmd.dpmodel.model.ener_model import ( + EnergyModel, +) + + +@BaseModel.register("dpa4_ener") +@BaseModel.register("sezm_ener") +class DPA4EnergyModel(EnergyModel): + r"""Energy model for the DPA4/SeZM descriptor family. + + Behaviorally identical to :class:`EnergyModel`; exists so the + dpa4-family wire types dispatch to a class that owns them instead of + polluting the generic energy model's registry. + """ diff --git a/deepmd/dpmodel/model/ener_model.py b/deepmd/dpmodel/model/ener_model.py index 09b50a6f17..a8280dbebf 100644 --- a/deepmd/dpmodel/model/ener_model.py +++ b/deepmd/dpmodel/model/ener_model.py @@ -36,8 +36,6 @@ @BaseModel.register("ener") -@BaseModel.register("sezm_ener") -@BaseModel.register("dpa4_ener") class EnergyModel(DPModelCommon, DPEnergyModel_): r"""Energy model that predicts total energy and derived quantities. diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index b7ff1a0293..44967bba74 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -281,6 +281,7 @@ def call_common( charge_spin: Array | None = None, neighbor_list: NeighborList | None = None, neighbor_graph_method: str | None = None, + spin: Array | None = None, ) -> dict[str, Array]: """Return model prediction. @@ -309,6 +310,23 @@ def call_common( The coordinates correction for virial. shape: nf x (nloc x 3) + charge_spin + Frame-level charge/spin FiLM conditioning, ``(nf, 2)`` or + ``None``. Both the dense (nlist) and NeighborGraph lowers + consume it (currently only DPA4/SeZM); the graph route no + longer forces this model onto dense (former ``cs -> dense`` + gate removed) -- it threads through + ``_call_common_graph``/``call_lower_graph`` to the + descriptor's ``call_graph``, gated per-descriptor by + ``supports_charge_spin``. + + spin + Per-local-atom spin, ``(nf, nloc, 3)``, or ``None``. Only the + NeighborGraph lower consumes it (native magnetic conditioning, + e.g. DPA4/SeZM); the dense (nlist) route has no spin support + and raises if ``spin`` is supplied without a graph + ``neighbor_graph_method``. + neighbor_list Neighbor-list construction strategy for the DENSE-nlist path only. ``None`` uses the default all-pairs builder; an @@ -350,10 +368,15 @@ def call_common( The keys are defined by the `ModelOutputDef`. """ - cc, bb, fp, ap, cs, input_prec = self._input_type_cast( - coord, box=box, fparam=fparam, aparam=aparam, charge_spin=charge_spin + cc, bb, fp, ap, cs, sp, input_prec = self._input_type_cast( + coord, + box=box, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + spin=spin, ) - del coord, box, fparam, aparam, charge_spin + del coord, box, fparam, aparam, charge_spin, spin graph_method = self._resolve_graph_method(neighbor_graph_method) # ``neighbor_list`` is a DENSE-nlist strategy; the graph path cannot # consume it. Reject an explicit graph+nlist combination, and @@ -367,10 +390,13 @@ def call_common( "pass one or the other" ) graph_method = None - # the graph lower does not consume charge_spin yet -> keep those - # models on dense (a None check, so it stays jit/export-safe) - if cs is not None: - graph_method = None + # model-level spin rides ONLY the NeighborGraph lower + if sp is not None and graph_method is None: + raise NotImplementedError( + "model-level spin rides only the NeighborGraph lower; the " + "dense (nlist) route has no spin support -- use a graph " + "neighbor_graph_method" + ) if graph_method is not None: # carry-all NeighborGraph energy forward (Option B / decision #17) model_predict = self._call_common_graph( @@ -381,6 +407,8 @@ def call_common( ap, graph_method, do_atomic_virial, + spin=sp, + charge_spin=cs, ) else: # legacy dense-nlist path (builds the extended quartet) @@ -445,6 +473,8 @@ def _call_common_graph( ap: Array | None, method: str, do_atomic_virial: bool = False, + spin: Array | None = None, + charge_spin: Array | None = None, ) -> dict[str, Array]: """Carry-all graph forward (opt-in, Option B). @@ -469,6 +499,16 @@ def _call_common_graph( the carry-all builder, ``"dense"`` or ``"ase"``. do_atomic_virial whether to calculate the atomic virial. + spin + Per-local-atom spin, ``(nf, nloc, 3)``, or ``None``. Flattened + to the flat node axis ``(N, 3)`` and forwarded unchanged to + :meth:`call_lower_graph`. + charge_spin + Frame-level charge/spin conditioning, ``(nf, 2)`` or ``None``. + Unflattened (per-frame, not per-node) and forwarded unchanged + to :meth:`call_lower_graph`, whose ``n_node`` here is always + the rectangular ``full(nf, nloc)`` this method builds -- the + one shape the descriptor's per-frame FiLM division requires. Returns ------- @@ -477,9 +517,7 @@ def _call_common_graph( (```` per-atom, ``_redu`` reduced, derivative name-holders ``None``, plus the int ``mask``). """ - descriptor = getattr(self.atomic_model, "descriptor", None) - uses_graph_lower = getattr(descriptor, "uses_graph_lower", lambda: False) - if not (self.mixed_types() and uses_graph_lower()): + if not (self.mixed_types() and self.atomic_model.uses_graph_lower()): raise NotImplementedError( "neighbor_graph_method requires a mixed_types descriptor with a " "graph lower (e.g. dpa1 attn_layer=0)" @@ -522,6 +560,8 @@ def _call_common_graph( if ap is not None else None ), + spin=(xp.reshape(spin, (nf * nloc, 3)) if spin is not None else None), + charge_spin=charge_spin, ) # Public ABI is rectangular (nf, nloc, *); the lower is flat # (N=nf*nloc, *). Unravel per-atom keys here at the boundary. @@ -594,7 +634,7 @@ def call_common_lower( nlist, extra_nlist_sort=self.need_sorted_nlist_for_lower(), ) - cc_ext, _, fp, ap, cs, input_prec = self._input_type_cast( + cc_ext, _, fp, ap, cs, _, input_prec = self._input_type_cast( extended_coord, fparam=fparam, aparam=aparam, charge_spin=charge_spin ) del extended_coord, fparam, aparam, charge_spin @@ -656,6 +696,7 @@ def forward_common_atomic_graph( aparam: Array | None = None, comm_dict: dict | None = None, charge_spin: Array | None = None, + spin: Array | None = None, ) -> dict[str, Array]: """Model-level graph forward (no type cast). Analogue of the dense :meth:`forward_common_atomic`. @@ -694,6 +735,10 @@ def forward_common_atomic_graph( Optional MPI communication metadata. charge_spin Charge/spin conditioning. + spin + Per-node spin vectors, flat (N, 3), or ``None``. Forwarded + unchanged to the atomic model's ``forward_common_atomic_graph`` + (and, from there, the descriptor's ``call_graph``). Returns ------- @@ -710,7 +755,12 @@ def forward_common_atomic_graph( n_local=n_local, ) atomic_ret = self.atomic_model.forward_common_atomic_graph( - graph, atype, fparam=fparam, aparam=aparam, charge_spin=charge_spin + graph, + atype, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + spin=spin, ) return fit_output_to_model_output_graph( atomic_ret, @@ -732,6 +782,7 @@ def call_common_lower_graph( aparam: Array | None = None, comm_dict: dict | None = None, charge_spin: Array | None = None, + spin: Array | None = None, ) -> dict[str, Array]: """Graph-native PUBLIC lower (dpa1/se_atten concat-tebd, attention included). @@ -769,14 +820,22 @@ def call_common_lower_graph( Optional MPI communication metadata. charge_spin Charge/spin conditioning. + spin + Per-node spin vectors, flat (N, 3), or ``None``. Cast to the + model precision alongside the other node inputs and forwarded + unchanged to :meth:`forward_common_atomic_graph`. Returns ------- dict The standard model dict in the INPUT precision. """ - edge_vec, _, fparam, aparam, cs, input_prec = self._input_type_cast( - edge_vec, fparam=fparam, aparam=aparam, charge_spin=charge_spin + edge_vec, _, fparam, aparam, cs, sp, input_prec = self._input_type_cast( + edge_vec, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + spin=spin, ) model_predict = self.forward_common_atomic_graph( atype, @@ -789,6 +848,7 @@ def call_common_lower_graph( aparam=aparam, comm_dict=comm_dict, charge_spin=cs, + spin=sp, ) model_predict = self._output_type_cast(model_predict, input_prec) return model_predict @@ -852,7 +912,16 @@ def _input_type_cast( fparam: Array | None = None, aparam: Array | None = None, charge_spin: Array | None = None, - ) -> tuple[Array, Array | None, Array | None, Array | None, Array | None, Any]: + spin: Array | None = None, + ) -> tuple[ + Array, + Array | None, + Array | None, + Array | None, + Array | None, + Array | None, + Any, + ]: """Cast the input data to global float type.""" xp = array_api_compat.array_namespace(coord) input_dtype = coord.dtype @@ -864,11 +933,11 @@ def _input_type_cast( ### _lst: list[Array | None] = [ xp.astype(vv, input_dtype) if vv is not None else None - for vv in [box, fparam, aparam, charge_spin] + for vv in [box, fparam, aparam, charge_spin, spin] ] - box, fparam, aparam, charge_spin = _lst + box, fparam, aparam, charge_spin, spin = _lst if input_dtype == global_dtype: - return coord, box, fparam, aparam, charge_spin, input_dtype + return coord, box, fparam, aparam, charge_spin, spin, input_dtype else: return ( xp.astype(coord, global_dtype), @@ -878,6 +947,7 @@ def _input_type_cast( xp.astype(charge_spin, global_dtype) if charge_spin is not None else None, + xp.astype(spin, global_dtype) if spin is not None else None, input_dtype, ) diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 8f96e965b0..bb3f846720 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -31,9 +31,15 @@ from deepmd.dpmodel.model.dp_zbl_model import ( DPZBLModel, ) +from deepmd.dpmodel.model.dpa4_model import ( + DPA4EnergyModel, +) from deepmd.dpmodel.model.ener_model import ( EnergyModel, ) +from deepmd.dpmodel.model.native_spin_model import ( + NativeSpinEnergyModel, +) from deepmd.dpmodel.model.polar_model import ( PolarModel, ) @@ -45,8 +51,11 @@ ) from deepmd.utils.spin import ( Spin, + normalize_spin_use_spin, ) +_DPA4_SEZM_DESCRIPTOR_TYPES = ("dpa4", "DPA4", "sezm", "SeZM") + def _get_standard_model_components( data: dict[str, Any], ntypes: int @@ -87,6 +96,14 @@ def get_standard_model(data: dict) -> EnergyModel: ) data = copy.deepcopy(data) ntypes = len(data["type_map"]) + # Analytical bridging (e.g. ZBL): the radii feed the DESCRIPTOR's + # InnerClamp/BridgingSwitch (mirrors pt's builder); the method builds the + # atomic model's InterPotential below. + bridging_method = str(data.get("bridging_method", "none")) + bridging_enabled = bridging_method.lower() not in ("none", "") + if bridging_enabled: + data["descriptor"]["inner_clamp_r_inner"] = data.get("bridging_r_inner", 0.5) + data["descriptor"]["inner_clamp_r_outer"] = data.get("bridging_r_outer", 0.8) descriptor, fitting, fitting_net_type = _get_standard_model_components(data, ntypes) atom_exclude_types = data.get("atom_exclude_types", []) pair_exclude_types = data.get("pair_exclude_types", []) @@ -99,6 +116,8 @@ def get_standard_model(data: dict) -> EnergyModel: modelcls = DOSModel elif fitting_net_type in ["ener", "direct_force_ener"]: modelcls = EnergyModel + elif fitting_net_type in ["dpa4_ener", "sezm_ener"]: + modelcls = DPA4EnergyModel elif fitting_net_type == "property": modelcls = PropertyModel else: @@ -111,6 +130,36 @@ def get_standard_model(data: dict) -> EnergyModel: atom_exclude_types=atom_exclude_types, pair_exclude_types=pair_exclude_types, ) + if bridging_enabled: + # Composition, not a flag (first-principles design): the analytical + # bridging term is its own atomic model, summed with the learned one + # by the existing linear composition machinery. + from deepmd.dpmodel.atomic_model.inter_potential import ( + InterPotentialAtomicModel, + ) + from deepmd.dpmodel.atomic_model.linear_atomic_model import ( + LinearEnergyAtomicModel, + ) + from deepmd.dpmodel.model.dp_linear_model import ( + LinearEnergyModel, + ) + + zbl_atomic = InterPotentialAtomicModel( + type_map=data["type_map"], + mode=bridging_method, + rcut=descriptor.get_rcut(), + sel=descriptor.get_sel(), + ) + composed = LinearEnergyAtomicModel( + models=[model.atomic_model, zbl_atomic], + type_map=data["type_map"], + weights="sum", + # Both exclusions belong to the composition: its children share one + # graph, so "excluded" must cover the analytical term too. + atom_exclude_types=atom_exclude_types, + pair_exclude_types=pair_exclude_types, + ) + return LinearEnergyModel(atomic_model_=composed) return model @@ -164,6 +213,11 @@ def get_spin_model(data: dict) -> SpinModel: data : dict The data to construct the model. """ + if data["descriptor"]["type"] in _DPA4_SEZM_DESCRIPTOR_TYPES: + raise NotImplementedError( + "the virtual-atom (deepspin) scheme is not supported for " + "DPA4/SeZM; use spin scheme 'native'" + ) data = copy.deepcopy(data) # include virtual spin and placeholder types data["type_map"] += [item + "_spin" for item in data["type_map"]] @@ -190,6 +244,69 @@ def get_spin_model(data: dict) -> SpinModel: return SpinModel(backbone_model=backbone_model, spin=spin) +def get_native_spin_model(data: dict) -> NativeSpinEnergyModel: + """Get a native (virtual-atom-free) spin model from a dictionary. + + Unlike :func:`get_spin_model`, no virtual atoms or doubled type map are + introduced: ``spin`` is injected into the descriptor config as + ``use_spin`` and consumed by the descriptor's equivariant spin + embedding. Any atomic model declaring ``supports_native_spin()`` is + eligible; the gate is that capability method, not a descriptor-type + list. + + The non-spin backbone is built by :func:`get_standard_model`, which OWNS + everything about assembling the atomic model -- descriptor/fitting, + exclusions and the analytical-bridging composition -- so ``spin`` and + ``bridging_method`` combine for free: the wrapper re-classes whatever + atomic model came back, be it a single learned model or a + ``LinearEnergyAtomicModel`` over ``[learned, InterPotential]`` (the + analytical child accepts and ignores ``spin``; the learned child consumes + it). + + Parameters + ---------- + data : dict + The data to construct the model. + """ + data = copy.deepcopy(data) + spin_cfg = data.pop("spin") + # Expand index/symbol forms of ``use_spin`` against ``type_map`` into the + # per-type boolean list (pure; validates symbols). + use_spin = normalize_spin_use_spin(spin_cfg["use_spin"], data["type_map"]) + spin = Spin( + use_spin=use_spin, + virtual_scale=spin_cfg.get("virtual_scale", 1.0), + allow_missing_label=spin_cfg.get("allow_missing_label", False), + ) + data.setdefault("descriptor", {}) + data["descriptor"]["use_spin"] = use_spin + try: + backbone_model = get_standard_model(data) + except TypeError as err: + if "use_spin" not in str(err): + # Unrelated construction error (e.g. a bogus fitting kwarg): + # propagate with its real context instead of masking it as a + # capability failure. + raise + # A descriptor without native spin support rejects the injected + # ``use_spin`` keyword at construction; translate to the + # capability-gate error. + raise NotImplementedError( + "spin scheme 'native' requires a descriptor with native spin " + "support (supports_native_spin()); descriptor type " + f"{data['descriptor'].get('type')!r} does not accept `use_spin`" + ) from err + # The ATOMIC MODEL answers the capability -- it knows its own structure, + # so this holds for a plain descriptor+fitting model and for a bridging + # composition alike, with no assumption here about either. + if not backbone_model.atomic_model.supports_native_spin(): + raise NotImplementedError( + "spin scheme 'native' requires an atomic model declaring " + "supports_native_spin()" + ) + return NativeSpinEnergyModel(atomic_model_=backbone_model.atomic_model, spin=spin) + + def get_model(data: dict) -> BaseModel: """Get a model from a dictionary. @@ -201,6 +318,8 @@ def get_model(data: dict) -> BaseModel: model_type = data.get("type", "standard") if model_type == "standard": if "spin" in data: + if data["spin"].get("scheme", "deepspin") == "native": + return get_native_spin_model(data) return get_spin_model(data) elif "use_srtab" in data: return get_zbl_model(data) diff --git a/deepmd/dpmodel/model/native_spin_model.py b/deepmd/dpmodel/model/native_spin_model.py new file mode 100644 index 0000000000..15943ba206 --- /dev/null +++ b/deepmd/dpmodel/model/native_spin_model.py @@ -0,0 +1,274 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Native-spin model factory (``make_native_spin_model``) and its concrete +energy-model instantiation (``NativeSpinEnergyModel``). +""" + +from copy import ( + deepcopy, +) +from typing import ( + Any, +) + +import numpy as np + +from deepmd.dpmodel.model.base_model import ( + BaseModel, +) +from deepmd.dpmodel.model.ener_model import ( + EnergyModel, +) +from deepmd.dpmodel.output_def import ( + ModelOutputDef, +) +from deepmd.utils.spin import ( + Spin, +) + + +class NativeSpinModelKind: + """Marker base identifying classes produced by ``make_native_spin_model``. + + Each backend instantiates the factory on its OWN standard model class, + so the concrete classes (e.g. dpmodel's and pt_expt's + ``NativeSpinEnergyModel``) are parallel products with NO subclass + relation between them -- an ``isinstance`` against one backend's + concrete class is silently dead in the other. Backend seams that need a + cross-backend family test (e.g. the with-comm freeze gate: native-spin + lowers are single-rank only) test against this shared marker instead. + """ + + +def make_native_spin_model(T_Model: type) -> type: + """Make a native-spin model class from a standard model class. + + The native scheme injects the per-atom spin vector directly into the + descriptor as an equivariant feature and obtains the magnetic force as + the negative spin gradient of the energy. No virtual atoms are created + (unlike :class:`~deepmd.dpmodel.model.spin_model.SpinModel`), so the + neighbor list, type map and selection stay at the real-system sizes. + + Mirrors :func:`~deepmd.dpmodel.model.make_model.make_model`'s + class-factory pattern: the produced class subclasses ``T_Model`` (is-a), + serializes as the parent's flat dict plus a ``spin`` field under wire + type ``"native_spin"``, and is meant to be registered in each backend's + ``BaseModel`` plugin registry so ``deserialize`` dispatch stays + backend-aware. Eligibility of a backbone is the + ``descriptor.supports_native_spin()`` capability, checked by the config + builders -- the factory itself is descriptor-agnostic. + + Parameters + ---------- + T_Model : type + The standard model class to derive from (e.g. the backend's + ``EnergyModel``). + + Returns + ------- + type + The derived native-spin model class. + """ + + class NSM(T_Model, NativeSpinModelKind): + """Native-spin variant of ``T_Model`` (see ``make_native_spin_model``).""" + + def __init__(self, *args: Any, spin: Spin, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.spin = spin + self.ntypes_real = self.spin.ntypes_real + # Per-real-type 0/1 spin gate. + self.spin_mask = self.spin.get_spin_mask() + + @staticmethod + def has_spin() -> bool: + """Returns whether it has spin input and output.""" + return True + + def model_output_def(self) -> ModelOutputDef: + """Get the spin-aware output def for the model.""" + atomic_output_def = self.atomic_output_def() + atomic_output_def["energy"].magnetic = True + return ModelOutputDef(atomic_output_def) + + def translated_output_def(self) -> dict[str, Any]: + """Get the translated output definition with public spin keys. + + Maps internal output names to user-facing names, e.g. + ``energy`` -> ``atom_energy``, ``energy_redu`` -> ``energy``, + ``energy_derv_r`` -> ``force``, ``energy_derv_r_mag`` -> + ``force_mag``. Built from this class's OWN + :meth:`model_output_def` (which sets ``energy.magnetic = True``). + """ + out_def_data = self.model_output_def().get_data() + model_output_type = self.model_output_type() + if "mask" in model_output_type: + model_output_type.pop(model_output_type.index("mask")) + var_name = model_output_type[0] + output_def = { + f"atom_{var_name}": out_def_data[var_name], + var_name: out_def_data[f"{var_name}_redu"], + "mask_mag": out_def_data["mask_mag"], + } + if self.do_grad_r(var_name): + output_def["force"] = deepcopy(out_def_data[f"{var_name}_derv_r"]) + output_def["force"].squeeze(-2) + output_def["force_mag"] = deepcopy( + out_def_data[f"{var_name}_derv_r_mag"] + ) + output_def["force_mag"].squeeze(-2) + if self.do_grad_c(var_name): + output_def["virial"] = deepcopy(out_def_data[f"{var_name}_derv_c_redu"]) + output_def["virial"].squeeze(-2) + output_def["atom_virial"] = deepcopy(out_def_data[f"{var_name}_derv_c"]) + output_def["atom_virial"].squeeze(-2) + return output_def + + def _spin_active_mask(self, atype: np.ndarray) -> np.ndarray: + """Single owner of ``mask_mag``: ``(N|nf,nloc, 1)`` bool, True + where the atom type carries spin. Array-api compatible; reused by + the eager and graph-export translations. + """ + return (self.spin_mask[atype] > 0)[..., None] + + def _translate_eager_call( + self, + model_ret: dict[str, np.ndarray], + atype: np.ndarray, + do_atomic_virial: bool = False, + ) -> dict[str, np.ndarray | None]: + """Single owner of the native-spin output translation, shared by + dpmodel/pt_expt ``call``/``forward`` and the graph export. Each + derivative key is the backend's value (pt_expt autograd) or + ``None`` (energy-only dpmodel); ``atom_virial`` is treated like + ``virial``, gated on ``do_atomic_virial``. + """ + out: dict[str, np.ndarray | None] = { + "atom_energy": model_ret["energy"], + "energy": model_ret["energy_redu"], + "mask_mag": self._spin_active_mask(atype), + } + translated = [ + ("energy_derv_r", "force"), + ("energy_derv_r_mag", "force_mag"), + ("energy_derv_c_redu", "virial"), + ] + if do_atomic_virial: + # Per-atom virial is opt-in (2.5x cost) -- only surfaced when + # requested, unlike the always-present reduced keys above. + translated.append(("energy_derv_c", "atom_virial")) + for kk_src, kk_dst in translated: + src = model_ret.get(kk_src) + out[kk_dst] = np.squeeze(src, axis=-2) if src is not None else None + return out + + def call( + self, + coord: np.ndarray, + atype: np.ndarray, + spin: np.ndarray, + box: np.ndarray | None = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + do_atomic_virial: bool = False, + charge_spin: np.ndarray | None = None, + ) -> dict[str, np.ndarray]: + """Return native-spin model predictions with translated public keys. + + Parameters + ---------- + coord + The coordinates of the atoms. shape: nf x (nloc x 3) + atype + The type of atoms. shape: nf x nloc + spin + The per-local-atom spin. shape: nf x (nloc x 3) + box + The simulation box. shape: nf x 9 + fparam + frame parameter. nf x ndf + aparam + atomic parameter. nf x nloc x nda + do_atomic_virial + If set, request the per-atom virial (``atom_virial`` key). + The energy-only dpmodel backend produces no derivatives, so + the value is ``None`` here (same as ``force``/``virial``); + the pt_expt subclass fills it with a real autograd tensor. + charge_spin + Frame-level charge/spin FiLM conditioning, shape + nf x dim_chg_spin (only consumed when the descriptor + declares ``add_chg_spin_ebd``). + + Returns + ------- + ret_dict + The result dict with translated keys: ``atom_energy``, + ``energy``, ``mask_mag``, plus + ``force``/``force_mag``/``virial`` (and ``atom_virial`` when + ``do_atomic_virial``) as ``None`` placeholders when the + backend produces no derivatives (dpmodel; the pt_expt + subclass produces real autograd tensors). + """ + model_ret = self.call_common( + coord, + atype, + box=box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + spin=spin, + charge_spin=charge_spin, + # dpmodel: opt into the carry-all NeighborGraph builder (the + # only lower that consumes model-level spin). + neighbor_graph_method="dense", + ) + return self._translate_eager_call( + model_ret, atype, do_atomic_virial=do_atomic_virial + ) + + def serialize(self) -> dict: + data = super().serialize() + # The backbone's own wire type would be lost under "native_spin"; + # keep it so deserialize can rebuild the RIGHT backbone. It is + # not always "standard": with analytical bridging the backbone is + # a composition ("linear"), whose dict has a different shape and + # @version. + data["backbone_type"] = data.get("type", "standard") + data["type"] = "native_spin" + data["spin"] = self.spin.serialize() + return data + + @classmethod + def deserialize(cls, data: dict) -> "NSM": + data = data.copy() + data.pop("type", None) + spin = Spin.deserialize(data.pop("spin")) + # make_model flat shape: the remaining dict IS the backbone + # (atomic) dict -- its @class/@version belong to the backbone's + # deserialize and must stay. Archives written before + # ``backbone_type`` existed are all plain standard models. + backbone_type = data.pop("backbone_type", "standard") + data["type"] = backbone_type + backbone_cls = ( + T_Model + if backbone_type == "standard" + else T_Model.get_class_by_type(backbone_type) + ) + backbone = backbone_cls.deserialize(data) + return cls(atomic_model_=backbone.atomic_model, spin=spin) + + return NSM + + +@BaseModel.register("native_spin") +class NativeSpinEnergyModel(make_native_spin_model(EnergyModel)): + r"""Native-spin energy model (dpmodel backend). + + dpmodel is energy-only for this model: it forwards through the + NeighborGraph lower (energy-only by design -- see + :meth:`~deepmd.dpmodel.model.make_model.make_model._call_common_graph`), + so ``call`` returns ``energy``/``atom_energy``/``mask_mag`` with + ``force``/``force_mag``/``virial`` as ``None`` placeholders. Force and + magnetic force are produced by autograd in the pt_expt backend. + Currently the DPA4/SeZM descriptor is the only one declaring + ``supports_native_spin()``. + """ diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index 0d9da07355..8db7abbbaf 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -71,7 +71,9 @@ def __init__( dp_atomic_model = self.backbone_model.get_dp_atomic_model() if dp_atomic_model is not None: descriptor = getattr(dp_atomic_model, "descriptor", None) - if descriptor is not None and hasattr(descriptor, "disable_graph_lower"): + if descriptor is not None: + # No-op on descriptors without a graph lower (BaseDescriptor + # concrete default). descriptor.disable_graph_lower() self.ntypes_real = self.spin.ntypes_real self.virtual_scale_mask = self.spin.get_virtual_scale_mask() diff --git a/deepmd/jax/jax_md/__init__.py b/deepmd/jax/jax_md/__init__.py index 13f23d212b..2d5b817199 100644 --- a/deepmd/jax/jax_md/__init__.py +++ b/deepmd/jax/jax_md/__init__.py @@ -267,7 +267,7 @@ def _normalize_fparam( if dim_fparam == 0: return None if fparam is None: - if getattr(model, "has_default_fparam", lambda: False)(): + if model.has_default_fparam(): default_fparam = model.get_default_fparam() if default_fparam is not None: return jnp.asarray(default_fparam, dtype=dtype).reshape(1, dim_fparam) diff --git a/deepmd/main.py b/deepmd/main.py index 03d1a2ec91..1206b2fa0e 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -362,9 +362,12 @@ def main_parser() -> argparse.ArgumentParser: type=str, choices=["nlist", "graph"], help="(Supported backend: PyTorch Exportable) Lower-level export form of the " - "frozen .pt2: 'nlist' (default, dense neighbor-list lower) or 'graph' " - "(NeighborGraph edge-list lower; only for graph-eligible models, currently " - "dpa1 with attn_layer=0). 'graph' selects the C++ graph inference path.", + "frozen .pt2: 'nlist' (dense neighbor-list lower) or 'graph' " + "(NeighborGraph edge-list lower, which selects the C++ graph inference " + "path). Only consulted for models that cannot use the graph lower: a " + "graph-capable model (DPA1/DPA2/DPA4, and every DPA model in future) " + "always freezes to 'graph', since the dense lower is deprecated in this " + "backend -- requesting 'nlist' for one logs a warning and is overridden.", ) # * test script ******************************************************************** diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py index c8bb113495..749bcdb76b 100644 --- a/deepmd/pd/infer/deep_eval.py +++ b/deepmd/pd/infer/deep_eval.py @@ -196,12 +196,10 @@ def __init__( else: raise TypeError("auto_batch_size should be bool, int, or AutoBatchSize") self._has_spin = ( - getattr(self.dp.model["Default"], "has_spin", False) + self.dp.model["Default"].has_spin() if isinstance(self.dp, ModelWrapper) else False ) - if callable(self._has_spin): - self._has_spin = False self._has_hessian = False def get_rcut(self) -> float: diff --git a/deepmd/pd/train/training.py b/deepmd/pd/train/training.py index 92225feeda..b2a8e6f5e5 100644 --- a/deepmd/pd/train/training.py +++ b/deepmd/pd/train/training.py @@ -1330,10 +1330,7 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: ) ] additional_data_requirement += aparam_requirement_items - has_spin = getattr(_model, "has_spin", False) - if callable(has_spin): - has_spin = has_spin() - if has_spin: + if _model.has_spin(): spin_requirement_items = [ DataRequirementItem("spin", ndof=3, atomic=True, must=True) ] diff --git a/deepmd/pd/train/wrapper.py b/deepmd/pd/train/wrapper.py index 05817920da..9db53a2156 100644 --- a/deepmd/pd/train/wrapper.py +++ b/deepmd/pd/train/wrapper.py @@ -165,10 +165,7 @@ def forward( "aparam": aparam, "charge_spin": charge_spin, } - has_spin = getattr(self.model[task_key], "has_spin", False) - if callable(has_spin): - has_spin = has_spin() - if has_spin: + if self.model[task_key].has_spin(): input_dict["spin"] = spin if self.inference_only or inference_only: diff --git a/deepmd/pt_expt/descriptor/dpa4.py b/deepmd/pt_expt/descriptor/dpa4.py index 65a4b94efd..5060de1a55 100644 --- a/deepmd/pt_expt/descriptor/dpa4.py +++ b/deepmd/pt_expt/descriptor/dpa4.py @@ -8,9 +8,11 @@ from deepmd.dpmodel.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4DP from deepmd.dpmodel.descriptor.dpa4_nn.activation import SwiGLU as SwiGLUDP from deepmd.dpmodel.descriptor.dpa4_nn.grid_net import GridProduct as GridProductDP +from deepmd.dpmodel.descriptor.dpa4_nn.radial import BridgingSwitch as BridgingSwitchDP from deepmd.dpmodel.descriptor.dpa4_nn.radial import ( C3CutoffEnvelope as C3CutoffEnvelopeDP, ) +from deepmd.dpmodel.descriptor.dpa4_nn.radial import InnerClamp as InnerClampDP from deepmd.kernels.utils import ( use_amp_infer, ) @@ -42,6 +44,32 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: return self.call(*args, **kwargs) +@torch_module +class InnerClamp(InnerClampDP): + def forward(self, *args: Any, **kwargs: Any) -> Any: + return self.call(*args, **kwargs) + + +# InnerClamp/BridgingSwitch are parameter-free (scalar bridging radii only, +# no serialize()); rebuild fresh from the stored constructor arguments. +register_dpmodel_mapping( + InnerClampDP, + lambda v: InnerClamp(v.r_inner, v.r_outer), +) + + +@torch_module +class BridgingSwitch(BridgingSwitchDP): + def forward(self, *args: Any, **kwargs: Any) -> Any: + return self.call(*args, **kwargs) + + +register_dpmodel_mapping( + BridgingSwitchDP, + lambda v: BridgingSwitch(v.r_inner, v.r_outer), +) + + # C3CutoffEnvelope carries only scalar configuration (cutoff radius and # polynomial exponent) and holds no trainable arrays, so it implements no # serialize()/deserialize() that the generic auto-wrap path relies on; rebuild @@ -164,6 +192,19 @@ class DescrptDPA4(DescrptDPA4DP): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) + # Persisted graph-routing knob (first-class training configuration): + # ``disable_graph_lower()`` used to flip only the plain dpmodel bool, + # which a Trainer checkpoint restart silently reset (the fresh model + # is rebuilt from config before ``load_state_dict``, and neither the + # state-dict keys nor ``_extra_state.model_params`` carried the + # choice) -- on a binding-sel system that switched the training + # equation and gradients without warning. A persistent buffer rides + # every pt_expt state_dict, so save/restart round-trips it. + torch.nn.Module.register_buffer( + self, + "graph_lower_disabled", + torch.zeros((), dtype=torch.bool, device="cpu"), + ) self.use_amp_infer = use_amp_infer() _promote_trainable_tree(self) @@ -174,6 +215,53 @@ def deserialize(cls, data: dict) -> "DescrptDPA4": obj = super().deserialize(data) return _promote_trainable_tree(obj) + def _in_training_mode(self) -> bool: + """Torch runtime hook for the training-only random local-Z roll. + + Overrides the dpmodel default (``False``) with the torch module's + ``training`` flag, restoring pt's ``random_gamma=self.random_gamma + and self.training`` semantics: train-mode forwards draw a fresh + gamma per call, eval/export forwards fix gamma (the export path + calls ``model.eval()`` before tracing). + """ + return bool(self.training) + + def disable_graph_lower(self) -> None: + """Persisted variant of the dpmodel escape hatch (see base class). + + The buffer (and the routing bool) are PER-TASK state: multi-task + ``share_params`` shares network submodules, not this buffer, so + disabling the graph lower on one task branch does not propagate to + branches sharing the same descriptor weights -- each branch owns + its routing decision. + """ + super().disable_graph_lower() + self.graph_lower_disabled.fill_(True) + + def _load_from_state_dict( + self, + state_dict: dict[str, Any], + prefix: str, + *args: Any, + **kwargs: Any, + ) -> None: + # Back-compat: checkpoints written before the knob was persisted lack + # the buffer; default to the fresh module's value (graph enabled) + # instead of failing the strict load. + key = prefix + "graph_lower_disabled" + if key not in state_dict: + state_dict[key] = self.graph_lower_disabled.detach().clone() + else: + # Re-sync the dpmodel-side routing bool from the RESTORED value + # here, at load time, where the incoming tensor is real. The + # routing predicate itself must stay a plain python bool: + # ``uses_graph_lower()`` runs inside traced forwards (the dense + # adapter gate), and reading the buffer there would emit a + # data-dependent ``bool(FakeTensor)`` guard that breaks + # torch.export (GuardOnDataDependentSymNode Eq(u0, 1)). + self._graph_lower_disabled = bool(state_dict[key]) + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + def forward(self, *args: Any, **kwargs: Any) -> Any: return self.call(*args, **kwargs) diff --git a/deepmd/pt_expt/descriptor/dpa4_nn/block.py b/deepmd/pt_expt/descriptor/dpa4_nn/block.py index c78196ff58..af14de1b49 100644 --- a/deepmd/pt_expt/descriptor/dpa4_nn/block.py +++ b/deepmd/pt_expt/descriptor/dpa4_nn/block.py @@ -29,22 +29,71 @@ from deepmd.dpmodel.descriptor.dpa4_nn.block import ( SeZMInteractionBlock as SeZMInteractionBlockDP, ) -from deepmd.dpmodel.descriptor.dpa4_nn.block import ( - exchange_ghost_features, -) from deepmd.pt_expt.common import ( torch_module, ) if TYPE_CHECKING: from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import ( - EdgeFeatureCache, + EdgeCache, ) # Environment values that enable an inference flag. _TRUTHY = {"1", "true", "yes", "on"} +def exchange_ghost_features( + x: torch.Tensor, + comm_dict: dict[str, torch.Tensor], +) -> torch.Tensor: + """Refresh ghost-node rows from their owner ranks via ``border_op``. + + Port of the pt-native SeZM exchange (``sezm_nn/block.py``): the node + multipole tensor is flattened to ``(nall, D*1*C)`` rows and exchanged + whole-row; SO(3) coefficients live in the shared global frame, so the + owner-to-ghost copy is exact and equivariant. ``border_op`` carries a + registered backward for gradient reverse-communication. + + Parameters + ---------- + x + Extended node features with shape ``(nall, D, 1, C)`` in block + precision. Owned rows lead; ghost rows are overwritten. + comm_dict + Border-exchange tensors ``send_list``, ``send_proc``, ``recv_proc``, + ``send_num``, ``recv_num``, ``communicator``, ``nlocal``, ``nghost``. + + Returns + ------- + torch.Tensor + ``x`` with ghost rows refreshed, same shape. + + Raises + ------ + NotImplementedError + When ``comm_dict`` carries ``has_spin`` — spin models do not route + the DPA4 graph lower. + """ + if "has_spin" in comm_dict: + raise NotImplementedError("spin models do not route the DPA4 graph lower") + n_nodes, ebed_dim, n_focus, channels = x.shape + # border_op exchanges whole rows by raw pointer arithmetic, so the + # buffer must be contiguous; a strided view would corrupt the exchange. + g1 = x.reshape(n_nodes, ebed_dim * n_focus * channels).contiguous() + g1 = torch.ops.deepmd_export.border_op( + comm_dict["send_list"], + comm_dict["send_proc"], + comm_dict["recv_proc"], + comm_dict["send_num"], + comm_dict["recv_num"], + g1, + comm_dict["communicator"], + comm_dict["nlocal"], + comm_dict["nghost"], + ) + return g1.reshape(n_nodes, ebed_dim, n_focus, channels) + + @torch_module class SeZMInteractionBlock(SeZMInteractionBlockDP): """SeZM interaction block with eval-time activation checkpointing.""" @@ -76,7 +125,7 @@ def _use_infer_activation_checkpoint(self, *tensors: torch.Tensor) -> bool: def _run_so2_unit( self, x: torch.Tensor, - edge_cache: EdgeFeatureCache, + edge_cache: EdgeCache, radial_feat: torch.Tensor, comm_dict: dict[str, torch.Tensor] | None = None, ) -> torch.Tensor: diff --git a/deepmd/pt_expt/descriptor/dpa4_nn/so2.py b/deepmd/pt_expt/descriptor/dpa4_nn/so2.py index 19e6b6aec9..3b72e0d0e7 100644 --- a/deepmd/pt_expt/descriptor/dpa4_nn/so2.py +++ b/deepmd/pt_expt/descriptor/dpa4_nn/so2.py @@ -47,7 +47,7 @@ if TYPE_CHECKING: from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import ( - EdgeFeatureCache, + EdgeCache, ) @@ -234,7 +234,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self._value_path = make_cute_value_path(self) def _rotate_to_local( - self, x: torch.Tensor, edge_cache: EdgeFeatureCache + self, x: torch.Tensor, edge_cache: EdgeCache ) -> tuple[torch.Tensor, torch.Tensor | None]: if self.use_triton_infer and not self.training: # ``self._rotate_to_local_fn`` was bound in ``__init__`` (the block @@ -248,7 +248,7 @@ def _rotate_to_local( return super()._rotate_to_local(x, edge_cache) def _rotate_back( - self, x_local: torch.Tensor, edge_cache: EdgeFeatureCache, n_edge: int + self, x_local: torch.Tensor, edge_cache: EdgeCache, n_edge: int ) -> torch.Tensor: if self.use_triton_infer and not self.training: Dt_full = edge_cache.Dt_full @@ -268,7 +268,7 @@ def _rotate_back( def _flash_aggregate( self, x_local_flash: torch.Tensor, - edge_cache: EdgeFeatureCache, + edge_cache: EdgeCache, attn_alpha: torch.Tensor, x_l0_node: torch.Tensor, n_node: int, diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index e49a73a518..6b394b3229 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -479,20 +479,48 @@ def train( ) +def _default_output_path(output: str, lower_kind: str) -> str: + """Default a suffixless frozen-model path from the RESOLVED lower kind. + + Parameters + ---------- + output : str + Requested output path. An explicit ``.pte`` / ``.pt2`` suffix is + preserved; any other (or missing) suffix is replaced. + lower_kind : str + The resolved lower kind (after ``freeze``'s native-spin + resolution): ``"graph"`` selects ``.pt2`` (an AOTI ``.pt2`` archive + is what the C++ graph path consumes), anything else ``.pte``. + + Returns + ------- + str + The output path with a definite ``.pte`` / ``.pt2`` suffix. + """ + if not output.endswith((".pte", ".pt2")): + return str( + Path(output).with_suffix(".pt2" if lower_kind == "graph" else ".pte") + ) + return output + + def freeze( model: str, - output: str = "frozen_model.pte", + output: str = "frozen_model", head: str | None = None, lower_kind: str = "nlist", ) -> None: - """Freeze a pt_expt checkpoint into a .pte exported model. + """Freeze a pt_expt checkpoint into a .pte/.pt2 exported model. Parameters ---------- model : str Path to the checkpoint file (.pt). output : str - Path for the output .pte file. + Path for the output file. When it carries no ``.pte``/``.pt2`` + suffix, the suffix is chosen from the RESOLVED lower kind (after the + native-spin resolution below): ``.pt2`` for a graph lower, ``.pte`` + otherwise. head : str or None Head to freeze in multi-task mode. lower_kind : str @@ -507,6 +535,10 @@ def freeze( softmax denominator, so graph-form results are sel-independent and differ from the legacy dense lower by up to ~1e-4 (see ``DescrptDPA1.call_graph``). + A NATIVE-spin model (``NativeSpinEnergyModel``) implements ONLY the + graph lower, so ``lower_kind`` is resolved to ``"graph"`` for it + regardless of the requested value -- before the export ABI and the + default output suffix are selected. """ import torch @@ -568,6 +600,35 @@ def freeze( m.eval() + # A graph-capable model ALWAYS freezes through the NeighborGraph lower. + # The dense lower is deprecated here and every DPA model is expected to be + # graph-capable, so "can use graph" is deliberately the whole condition -- + # a separate "graph-required" capability would encode a distinction that + # stops existing. It is also load-bearing: for native spin and for an + # analytical bridging term no dense lower exists at all, so the public + # default would otherwise fail deep inside the dense trace. Resolved HERE, + # before the export ABI and the default output suffix are chosen. + from deepmd.pt_expt.model.graph_lower import ( + model_uses_graph_lower, + ) + + if lower_kind != "graph" and model_uses_graph_lower(m): + # WARNING, not info: this overrides what the caller asked for, and + # the two lowers are not numerically identical (dpa1's graph and + # dense attention semantics differ by ~1e-4), so the override must + # be visible in the log rather than inferred from the output suffix. + log.warning( + "Requested lower_kind=%r is being OVERRIDDEN to 'graph': the " + "dense (nlist) lower is deprecated in the pt_expt backend and " + "this model is graph-lower capable. The frozen artifact will be " + "a graph-kind .pt2, and its outputs may differ slightly from the " + "dense lower.", + lower_kind, + ) + lower_kind = "graph" + + output = _default_output_path(output, lower_kind) + # The graph lower is opt-in and only valid for graph-eligible models # (dpa1 with concat tebd, incl. attention layers and exclude_types # -- the carry-all pair enumeration exports via unbacked SymInts). Fail @@ -833,18 +894,14 @@ def main(args: list[str] | argparse.Namespace | None = None) -> None: f"Checkpoint path '{model_path}' does not exist." ) FLAGS.model = str(model_path) - _lower_kind = getattr(FLAGS, "lower_kind", "nlist") - if not FLAGS.output.endswith((".pte", ".pt2")): - # Default suffix: .pt2 for the graph export (an AOTI .pt2 archive is - # what the C++ graph path consumes), .pte otherwise. Explicit user - # .pte / .pt2 suffixes are preserved for both. - _default_suffix = ".pt2" if _lower_kind == "graph" else ".pte" - FLAGS.output = str(Path(FLAGS.output).with_suffix(_default_suffix)) + # Suffix defaulting lives in freeze(): the correct suffix depends on + # the RESOLVED lower kind (native-spin models force 'graph'), which + # is only known after the checkpoint's model is built there. freeze( model=FLAGS.model, output=FLAGS.output, head=FLAGS.head, - lower_kind=_lower_kind, + lower_kind=getattr(FLAGS, "lower_kind", "nlist"), ) elif FLAGS.command == "change-bias": change_bias( diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 1b97ef7f50..999e7f69f1 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json +import warnings from collections.abc import ( Callable, ) @@ -90,6 +91,34 @@ } +def _graph_spin_output_key(odef: "OutputVariableDef") -> str | None: + """Map a native-spin request def to its graph-spin public model key. + + Twin of ``_GRAPH_CATEGORY_TO_KEY`` for + ``NativeSpinEnergyModel.forward_lower_graph_exportable``'s output dict + (``atom_energy``, ``energy``, ``force``, ``force_mag``, ``virial``, + ``mask_mag``, ``atom_virial``). Category alone + is NOT enough here: ``do_derivative`` (``deepmd.dpmodel.output_def``) + gives the magnetic derivative def (``energy_derv_r_mag``) the SAME + category as the physical one (``energy_derv_r``) -- only ``.magnetic`` + tells them apart -- so this checks ``magnetic`` before falling back to + the category table. ``mask_mag`` is exported directly by the graph + forward (the model owns the derivation -- see + ``NativeSpinEnergyModel._spin_active_mask``), so it maps to its own key + and is read like any other output. Returns ``None`` only for ``mask`` + (the real-atom mask, a virtual-atom-scheme concept the native-spin + graph ABI does not produce); that falls back to the NaN-filled + placeholder, same as any other unavailable output. + """ + if odef.name == "mask": + return None + if odef.name == "mask_mag": + return "mask_mag" + if odef.magnetic and odef.category == OutputVariableCategory.DERV_R: + return "force_mag" + return _GRAPH_CATEGORY_TO_KEY.get(odef.category) + + def _reshape_charge_spin( charge_spin: np.ndarray, nframes: int, dim_chg_spin: int ) -> np.ndarray: @@ -115,6 +144,29 @@ def _is_pt_backend_dpa4_params(model_params: dict[str, Any]) -> bool: return False +def _warn_legacy_edge_vec(metadata: dict) -> None: + """Warn once per model load when an edge_vec-schema artifact is opened. + + The ``edge_vec`` lower schema is produced only by the pt backend's + SeZM/DPA4 freeze and is superseded by the NeighborGraph lower. Support + will be removed in a future release; energy SeZM checkpoints can be + refrozen through the pt_expt backend (graph schema) instead. + + Parameters + ---------- + metadata + The ``metadata.json`` dict of the opened ``.pt2`` archive. + """ + if metadata.get("lower_input_kind") == "edge_vec": + warnings.warn( + "This .pt2 uses the deprecated edge_vec lower schema (pt-backend " + "SeZM/DPA4 freeze). Support will be removed in a future release; " + "refreeze the checkpoint with the pt_expt backend (graph schema).", + DeprecationWarning, + stacklevel=3, + ) + + class DeepEval(DeepEvalBackend): """PyTorch Exportable backend implementation of DeepEval. @@ -184,6 +236,11 @@ def __init__( "`.pt` (training checkpoint)." ) + # Single choke point: self.metadata is set by all three loaders + # above (_load_pt2 / _load_pte / _load_pt), so this is the one + # place that sees every model load regardless of archive kind. + _warn_legacy_edge_vec(self.metadata) + # neighbor_graph_method is consumed ONLY by graph-form .pt2 eval # (_eval_model_graph); fail fast instead of silently ignoring it on # nlist-form artifacts (there, the builder knob is nlist_backend). @@ -226,9 +283,17 @@ def _setup_nlist_backend(self, nlist_backend: str) -> None: "expected 'auto', 'vesin', or 'native'." ) is_spin = bool(getattr(self, "_is_spin", False)) + # Native-spin (NeighborGraph route) graph-form artifacts never touch + # this NLIST builder at all -- graph-form eval uses + # ``neighbor_graph_method`` instead (see ``_eval_model_graph_spin``). + # Only the virtual-atom (dense/nlist) spin scheme actually needs the + # vesin restriction below. + is_native_spin_graph = is_spin and ( + getattr(self, "metadata", {}).get("lower_input_kind") == "graph" + ) 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 + unsupported = "spin models" if is_spin and not is_native_spin_graph else None if nlist_backend == "native": self._use_vesin = False elif nlist_backend == "vesin": @@ -268,7 +333,8 @@ def _init_from_model_json(self, model_json_str: str) -> None: model_dict = _json_to_numpy(model_dict) model_data = model_dict["model"] - if model_data.get("type") == "spin_ener": + model_type = model_data.get("type") + if model_type == "spin_ener": from deepmd.pt_expt.model.spin_model import ( SpinModel, ) @@ -276,8 +342,12 @@ def _init_from_model_json(self, model_json_str: str) -> None: self._dpmodel = SpinModel.deserialize(model_data) self._is_spin = True else: + # Registry-dispatched: wrapper classes registered in the pt_expt + # BaseModel registry (e.g. the native-spin models, type + # "native_spin") come back as their pt_expt torch classes and + # declare spin via the base-model capability method. self._dpmodel = BaseModel.deserialize(model_data) - self._is_spin = False + self._is_spin = self._dpmodel.has_spin() self._rcut = self._dpmodel.get_rcut() self._type_map = self._dpmodel.get_type_map() @@ -1584,6 +1654,22 @@ def _eval_model_spin( request_defs: list[OutputVariableDef], charge_spin: np.ndarray | None = None, ) -> tuple[np.ndarray, ...]: + if self.metadata.get("lower_input_kind") == "graph": + # Native-spin (NeighborGraph route): no virtual atoms and no + # extended/nlist ABI at all -- dispatch to the graph-native fast + # path (mirrors _eval_model's dispatch to _eval_model_graph for + # the non-spin case). charge_spin rides the conditional slot-13 + # tail (see NativeSpinEnergyModel.forward_lower_graph_exportable). + return self._eval_model_graph_spin( + coords, + cells, + atom_types, + spins, + fparam, + aparam, + request_defs, + charge_spin=charge_spin, + ) nframes = coords.shape[0] if len(atom_types.shape) == 1: natoms = len(atom_types) @@ -1751,6 +1837,136 @@ def _eval_model_spin( ) return tuple(results) + def _eval_model_graph_spin( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + spins: np.ndarray, + fparam: np.ndarray | None, + aparam: np.ndarray | None, + request_defs: list[OutputVariableDef], + charge_spin: np.ndarray | None = None, + ) -> tuple[np.ndarray, ...]: + """Evaluate a graph-form native-spin ``.pt2`` (``lower_input_kind == + "graph"`` and ``is_spin``). + + Mirrors :meth:`_eval_model_graph`'s carry-all + :class:`~deepmd.dpmodel.utils.neighbor_graph.NeighborGraph` + construction (SAME builder, SAME positional ABI up through + ``source_row_ptr``), then inserts the owned-atom ``spin`` tensor + ``(N, 3)`` at positional index 10 of + ``NativeSpinEnergyModel.forward_lower_graph_exportable`` -- the node + axis IS the owned-local-atom axis for single-rank eval (no ghost + nodes), so ``spin`` needs no extension/mapping, unlike the dense + spin path's ``ext_spin_t``. ``charge_spin`` rides the conditional + slot-13 tail (combined native-spin + charge-spin FiLM models; the + slot is dropped from the exported signature otherwise). The forward + returns LOCAL public keys directly (``atom_energy``, + ``energy``, ``force``, ``force_mag``, ``virial``, ``atom_virial``), + so results are reshaped without ``communicate_extended_output``, + same as the non-spin graph path. + """ + from deepmd.pt_expt.utils.env import ( + DEVICE, + ) + + 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]) + + coord_input = coords.reshape(nframes, natoms, 3) + box_input = cells.reshape(nframes, 9) if cells is not None else None + graph = self._build_eval_graph(coord_input, atom_types, box_input, DEVICE) + + atype_t = torch.tensor( + np.asarray(atom_types).reshape(-1), dtype=torch.int64, device=DEVICE + ) + n_node_t = torch.as_tensor(graph.n_node, dtype=torch.int64, device=DEVICE) + edge_index_t = torch.as_tensor( + graph.edge_index, dtype=torch.int64, device=DEVICE + ) + edge_dtype = ( + torch.float32 + if self.metadata.get("graph_edge_dtype") == "float32" + else torch.float64 + ) + edge_vec_t = torch.as_tensor( + graph.edge_vec, + dtype=edge_dtype, + device=DEVICE, + ) + edge_mask_t = torch.as_tensor(graph.edge_mask, dtype=torch.bool, device=DEVICE) + destination_order_t = torch.as_tensor( + graph.destination_order, + device=DEVICE, + ) + destination_row_ptr_t = torch.as_tensor( + graph.destination_row_ptr, + dtype=torch.int64, + device=DEVICE, + ) + source_order_t = torch.as_tensor( + graph.source_order, + device=DEVICE, + ) + source_row_ptr_t = torch.as_tensor( + graph.source_row_ptr, + dtype=torch.int64, + device=DEVICE, + ) + + spin_t = torch.tensor( + np.asarray(spins).reshape(nframes * natoms, 3), + dtype=torch.float64, + device=DEVICE, + ) + + fparam_t, aparam_t = self._prepare_optional_lower_inputs( + fparam, aparam, nframes, natoms, DEVICE + ) + if aparam_t is not None: + # graph-lower ABI: aparam is FLAT on the node axis, (N, nda) -- + # the same axis as ``atype``/``spin`` (mirrors _eval_model_graph). + aparam_t = aparam_t.reshape(nframes * natoms, -1) + + model_inputs = ( + atype_t, + n_node_t, + n_node_t, + edge_index_t, + edge_vec_t, + edge_mask_t, + destination_order_t, + destination_row_ptr_t, + source_order_t, + source_row_ptr_t, + spin_t, + fparam_t, + aparam_t, + self._make_charge_spin_input(nframes, charge_spin), + ) + if self._is_pt2: + model_ret = self._pt2_runner(*model_inputs) + else: + model_ret = self.exported_module(*model_inputs) + + results = [] + for odef in request_defs: + shape = self._get_output_shape(odef, nframes, natoms) + gkey = _graph_spin_output_key(odef) + val = model_ret.get(gkey) if gkey is not None else None + if val is not None: + results.append(val.detach().cpu().numpy().reshape(shape)) + else: + results.append( + np.full(np.abs(shape), np.nan, dtype=GLOBAL_NP_FLOAT_PRECISION) + ) + return tuple(results) + def _eval_model_graph( self, coords: np.ndarray, @@ -2144,9 +2360,14 @@ def eval_typeebd(self) -> np.ndarray: self._require_dpmodel("eval_typeebd") from deepmd.dpmodel.utils.type_embed import TypeEmbedNet as TypeEmbedNetDP + from deepmd.pt_expt.model.spin_model import ( + SpinModel, + ) model = self._dpmodel - if self._is_spin_model(): + if isinstance(model, SpinModel): + # Virtual-atom wrapper: type-embed nets live on the backbone. + # Native-spin models ARE the model (is-a); no unwrap. model = model.backbone_model out = [] for mm in model.modules(): diff --git a/deepmd/pt_expt/model/__init__.py b/deepmd/pt_expt/model/__init__.py index 9e9a24f0a8..b0d7ca1402 100644 --- a/deepmd/pt_expt/model/__init__.py +++ b/deepmd/pt_expt/model/__init__.py @@ -15,6 +15,9 @@ from .dp_zbl_model import ( DPZBLModel, ) +from .dpa4_model import ( + DPA4EnergyModel, +) from .ener_model import ( EnergyModel, ) @@ -27,6 +30,9 @@ from .model import ( BaseModel, ) +from .native_spin_model import ( + NativeSpinEnergyModel, +) from .polar_model import ( PolarModel, ) @@ -40,11 +46,13 @@ __all__ = [ "BaseModel", "DOSModel", + "DPA4EnergyModel", "DPZBLModel", "DipoleModel", "EnergyModel", "FrozenModel", "LinearEnergyModel", + "NativeSpinEnergyModel", "PolarModel", "PropertyModel", "SpinEnergyModel", diff --git a/deepmd/pt_expt/model/dp_linear_model.py b/deepmd/pt_expt/model/dp_linear_model.py index 4a29251932..e4d9e0b877 100644 --- a/deepmd/pt_expt/model/dp_linear_model.py +++ b/deepmd/pt_expt/model/dp_linear_model.py @@ -19,6 +19,9 @@ DeepmdDataSystem, ) +from .ener_model import ( + EnergyModel, +) from .make_model import ( _pad_nlist_for_export, make_model, @@ -30,8 +33,15 @@ DPLinearModel_ = make_model(LinearEnergyAtomicModel, T_Bases=(BaseModel,)) -@BaseModel.register("linear_ener") +@BaseModel.register("linear_ener") # config type +@BaseModel.register("linear") # wire type emitted by the flat serialize class LinearEnergyModel(DPModelCommon, DPLinearModel_): + # The graph .pt2 exportable is energy-contract machinery (public-key + # translation over the CM's forward_common_lower_graph_exportable), + # identical for any energy model -- reuse EnergyModel's verbatim so + # compositions (e.g. analytical bridging) freeze like standard models. + forward_lower_graph_exportable = EnergyModel.forward_lower_graph_exportable + def __init__( self, *args: Any, diff --git a/deepmd/pt_expt/model/dpa4_model.py b/deepmd/pt_expt/model/dpa4_model.py new file mode 100644 index 0000000000..2cc7355ad8 --- /dev/null +++ b/deepmd/pt_expt/model/dpa4_model.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""DPA4/SeZM-family energy model (pt_expt backend). + +A THIN subclass of the descriptor-agnostic :class:`EnergyModel` that owns +exactly the dpa4-family concerns: + +- the family's registry wire types (``dpa4_ener``/``sezm_ener`` + fitting-type dispatch keys and the pt model-type strings + ``dpa4``/``sezm``/``sezm_atomic``); +- the pt-checkpoint conversion (pt's ``SeZMModel`` wrapper and + ``sezm_atomic`` dict layouts), moved here from ``BaseModel`` so the + generic base assumes nothing about a specific descriptor family. + +May later absorb the dpa4-family builder (``get_sezm_model``) as a +``get_model`` classmethod, unifying dpa4 construction with the standard +registry dispatch. +""" + +from typing import ( + Any, +) + +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .ener_model import ( + EnergyModel, +) +from .model import ( + BaseModel, +) + + +@BaseModel.register("dpa4_ener") +@BaseModel.register("sezm_ener") +@BaseModel.register("dpa4") +@BaseModel.register("DPA4") +@BaseModel.register("sezm") +@BaseModel.register("SeZM") +@BaseModel.register("sezm_atomic") +class DPA4EnergyModel(EnergyModel): + r"""Energy model for the DPA4/SeZM descriptor family. + + Behaviorally identical to :class:`EnergyModel`; additionally owns the + pt-checkpoint interop: pt's ``SeZMModel`` serialises with a model-level + wrapper (``{type: "SeZM", atomic_model: , + bridging_method, bridging_r_*, lora}``) whose atomic dict carries + pt-only extras -- :meth:`deserialize` recognises those layouts, + normalises them to the standard flat dict, and delegates to the + generic path. + """ + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> "DPA4EnergyModel": + model_type = str(data.get("type", "standard")).lower() + if model_type in ("sezm", "dpa4"): + return cls.deserialize(cls._unwrap_pt_sezm_model(data)) + if model_type == "sezm_atomic": + return cls.deserialize(cls._normalize_pt_sezm_atomic(data)) + return super().deserialize(data) + + @staticmethod + def _unwrap_pt_sezm_model(data: dict[str, Any]) -> dict[str, Any]: + """Unwrap pt's ``SeZMModel`` serialization to the inner atomic dict.""" + data = data.copy() + # The pt SeZM model wrapper serialises with ``@version`` 1. Validate + # before discarding it so a future incompatible wrapper schema is not + # silently mis-deserialized (the wrapper only carries the guarded + # bridging/lora extras below, so the accepted range is narrow). + check_version_compatibility(int(data.get("@version", 1)), 1, 1) + bridging_method = str(data.get("bridging_method", "none")).lower() + if data.get("lora") is not None: + raise NotImplementedError( + "Deserializing a pt SeZM/DPA4 checkpoint with `lora` is " + "not supported in pt_expt." + ) + atomic_model = data.get("atomic_model") + if atomic_model is None: + raise ValueError( + "SeZM/DPA4 model data is missing the 'atomic_model' entry." + ) + if bridging_method not in ("none", ""): + # pt serializes bridging as a flag on its model wrapper; our + # architecture is a linear COMPOSITION with a different dict + # shape. No conversion is claimed -- rebuild from the training + # config instead. + raise NotImplementedError( + "Deserializing a pt SeZM/DPA4 checkpoint with " + f"`bridging_method`={data.get('bridging_method')!r} is not " + "supported; rebuild the bridged model from its config." + ) + return atomic_model + + @staticmethod + def _normalize_pt_sezm_atomic(data: dict[str, Any]) -> dict[str, Any]: + """Convert a pt ``sezm_atomic`` dict to a standard atomic dict. + + Strips the pt-only ``dens`` head state (``dens_fitting`` / + ``active_mode`` / the ``dens_force_rmsd`` @variable) and rewrites the + ``type``/``@version`` so the generic dpmodel atomic-model deserialize + accepts it. A non-energy active mode or a populated dens head is + rejected because pt_expt only implements the energy path. + """ + data = data.copy() + # pt emits ``@version`` 3 for ``sezm_atomic``; the standard dpmodel + # atomic-model deserialize requires exactly 2. The only schema delta + # between the two is the stripped ``dens`` state below, so coercion is + # safe for the known-compatible range {2, 3}. Validate the incoming + # version BEFORE coercing so a future incompatible pt schema (e.g. + # ``@version`` 4) is rejected loudly instead of mis-deserialized. + check_version_compatibility(int(data.get("@version", 2)), 3, 2) + if data.pop("dens_fitting", None) is not None: + raise NotImplementedError( + "Deserializing a pt SeZM/DPA4 checkpoint with a `dens` " + "fitting head is not supported in pt_expt." + ) + active_mode = data.pop("active_mode", None) + if active_mode not in (None, "ener"): + raise NotImplementedError( + f"Deserializing a pt SeZM/DPA4 checkpoint in active_mode " + f"{active_mode!r} is not supported in pt_expt (energy only)." + ) + variables = data.get("@variables") + if isinstance(variables, dict): + data["@variables"] = { + k: v for k, v in variables.items() if k in ("out_bias", "out_std") + } + # The standard dpmodel atomic-model deserialize checks @version == 2. + data["@version"] = 2 + data["type"] = "standard" + return data diff --git a/deepmd/pt_expt/model/edge_transform_output.py b/deepmd/pt_expt/model/edge_transform_output.py index 52761db645..7838781889 100644 --- a/deepmd/pt_expt/model/edge_transform_output.py +++ b/deepmd/pt_expt/model/edge_transform_output.py @@ -15,6 +15,9 @@ from deepmd.dpmodel.model.edge_transform_output import ( node_ownership_mask, ) +from deepmd.dpmodel.output_def import ( + get_deriv_name_mag, +) from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, edge_force_virial, @@ -156,6 +159,7 @@ def fit_output_to_model_output_graph( node_capacity: int | None = None, n_local: torch.Tensor | None = None, force_precision: torch.dtype | None = None, + spin_leaf: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Graph analogue of the dense pt_expt ``fit_output_to_model_output``. @@ -216,6 +220,17 @@ def fit_output_to_model_output_graph( Compute precision (model dtype) in which to assemble the force / virial during inference, decoupled from the fp64 ``edge_vec`` leaf; see :func:`edge_energy_deriv`. ``None`` keeps the leaf dtype. + spin_leaf + Per-node native spin autograd leaf, flat ``(N, 3)``, or ``None``. When + given, every ``r_differentiable`` reducible output additionally emits + ``_derv_r_mag = -d_redu/dspin`` (``(N, *shape, 3)``), via a + SECOND ``torch.autograd.grad`` call on the SAME reduced scalar that + the force grad differentiates (same energy, same backward family). + This deliberately differs from the pt backend + (``deepmd/pt/model/model/transform_output.py:288``), which computes + ``grad(E, [edge_vec, spin])`` jointly in one call; here the two grads + are separate so :func:`edge_energy_deriv`'s signature stays + untouched. ``None`` (default): unchanged, no mag output. Returns ------- @@ -295,6 +310,7 @@ def fit_output_to_model_output_graph( ff_list: list[torch.Tensor] = [] av_list: list[torch.Tensor] = [] vir_list: list[torch.Tensor] = [] + mag_list: list[torch.Tensor] = [] for c in range(size): force, atom_vir, vir = edge_energy_deriv( svv[:, c], @@ -320,8 +336,26 @@ def fit_output_to_model_output_graph( assert atom_vir is not None # atom_virial (N, 3, 3) -> (N, 1, 9) [flat] av_list.append(atom_vir.reshape(N, 1, 9)) + if spin_leaf is not None: + # Second, separate backward on the SAME reduced scalar + # ``svv[:, c]`` the force grad above just differentiated -- + # same energy, same backward family, different leaf. See the + # ``spin_leaf`` docstring entry for the pt-parity note. + (g_s,) = torch.autograd.grad( + svv[:, c].sum(), + spin_leaf, + create_graph=create_graph, + retain_graph=True, + ) + # force_mag (N, 3) -> (N, 1, 3) [flat; caller unravels] + mag_list.append((-g_s).reshape(N, 1, 3)) # (N, size, 3) -> (N, *shape, 3) model_ret[kk_derv_r] = torch.cat(ff_list, dim=-2).reshape([N, *shap, 3]) + if spin_leaf is not None: + kk_derv_r_mag, _ = get_deriv_name_mag(kk) + model_ret[kk_derv_r_mag] = torch.cat(mag_list, dim=-2).reshape( + [N, *shap, 3] + ) if vdef.c_differentiable: # (nf, size, 9) -> (nf, *shape, 9) model_ret[kk_derv_c + "_redu"] = torch.cat(vir_list, dim=-2).reshape( diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index 390afba79b..b5b650ab40 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -68,8 +68,6 @@ def _translate_energy_keys( @BaseModel.register("ener") -@BaseModel.register("sezm_ener") -@BaseModel.register("dpa4_ener") class EnergyModel(DPModelCommon, DPEnergyModel_): def __init__( self, @@ -168,7 +166,10 @@ def forward_lower_canonical_graph( fitting, graph, atype, - descriptor.type_embedding.call(), + # descriptor-owned hook (single owner for the graph-route tebd + # table); value-identical for dpa1, the only canonical-eligible + # descriptor. + descriptor.graph_type_embedding_table(), output_mask, atom_bias, do_atomic_virial, diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index 7efa904f23..8781f3bb22 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -26,12 +26,18 @@ from deepmd.pt_expt.model.dos_model import ( DOSModel, ) +from deepmd.pt_expt.model.dpa4_model import ( + DPA4EnergyModel, +) from deepmd.pt_expt.model.ener_model import ( EnergyModel, ) from deepmd.pt_expt.model.model import ( BaseModel, ) +from deepmd.pt_expt.model.native_spin_model import ( + NativeSpinEnergyModel, +) from deepmd.pt_expt.model.polar_model import ( PolarModel, ) @@ -43,6 +49,7 @@ ) from deepmd.utils.spin import ( Spin, + normalize_spin_use_spin, ) log = logging.getLogger(__name__) @@ -141,13 +148,18 @@ def get_sezm_model(data: dict) -> EnergyModel: ) _WARNED_ONCE.add("enable_tf32") if "spin" in data: - raise NotImplementedError( - "Spin DPA4/SeZM models are not supported in the pt_expt backend." - ) - if str(data.get("bridging_method", "none")).lower() != "none": - raise NotImplementedError( - "`bridging_method` is not supported for DPA4/SeZM in the pt_expt backend." - ) + if str(data["spin"].get("scheme", "deepspin")) != "native": + raise NotImplementedError( + "Spin DPA4/SeZM models with the virtual-atom (deepspin) " + "scheme are not supported in the pt_expt backend; use spin " + "scheme 'native' instead." + ) + return get_native_spin_model(data) + # Analytical bridging (e.g. ZBL): the radii feed the DESCRIPTOR's + # InnerClamp/BridgingSwitch (mirrors pt's builder); the method builds the + # atomic model's InterPotential at construction below. + bridging_method = str(data.get("bridging_method", "none")) + bridging_enabled = bridging_method.lower() not in ("none", "") if data.get("lora") is not None: raise NotImplementedError( "`lora` is not supported for DPA4/SeZM in the pt_expt backend." @@ -163,6 +175,9 @@ def get_sezm_model(data: dict) -> EnergyModel: data.pop("type", None) data.setdefault("descriptor", {}) data.setdefault("fitting_net", {}) + if bridging_enabled: + data["descriptor"]["inner_clamp_r_inner"] = data.get("bridging_r_inner", 0.5) + data["descriptor"]["inner_clamp_r_outer"] = data.get("bridging_r_outer", 0.8) data["descriptor"].setdefault("type", "dpa4") data["fitting_net"].setdefault("type", "dpa4_ener") # the DPA4/SeZM model type is a fixed descriptor/fitting contract; reject @@ -196,13 +211,108 @@ def get_sezm_model(data: dict) -> EnergyModel: ntypes = len(data["type_map"]) descriptor, fitting, _ = _get_standard_model_components(data, ntypes) - return EnergyModel( + model = DPA4EnergyModel( descriptor=descriptor, fitting=fitting, type_map=data["type_map"], atom_exclude_types=data.get("atom_exclude_types", []), pair_exclude_types=pair_exclude_types, ) + if bridging_enabled: + # Composition, not a flag (first-principles design): the analytical + # bridging term is its own atomic model, summed with the learned one + # by the existing linear composition machinery. + from deepmd.dpmodel.atomic_model.inter_potential import ( + InterPotentialAtomicModel, + ) + from deepmd.dpmodel.atomic_model.linear_atomic_model import ( + LinearEnergyAtomicModel, + ) + from deepmd.pt_expt.model.dp_linear_model import ( + LinearEnergyModel, + ) + + zbl_atomic = InterPotentialAtomicModel( + type_map=data["type_map"], + mode=bridging_method, + rcut=descriptor.get_rcut(), + sel=descriptor.get_sel(), + ) + composed = LinearEnergyAtomicModel( + models=[model.atomic_model, zbl_atomic], + type_map=data["type_map"], + weights="sum", + # Both exclusions belong to the composition: its children share one + # graph, so "excluded" must cover the analytical term too. + atom_exclude_types=data.get("atom_exclude_types", []), + pair_exclude_types=pair_exclude_types, + ) + return LinearEnergyModel(atomic_model_=composed) + return model + + +def get_native_spin_model(data: dict) -> NativeSpinEnergyModel: + """Build a pt_expt native (virtual-atom-free) spin model. + + Mirrors :func:`deepmd.dpmodel.model.model.get_native_spin_model`: no + virtual atoms or doubled type map are introduced, and ``use_spin`` is + injected into the descriptor config (consumed by the descriptor's + equivariant spin embedding). The non-spin backbone is built by the + standard builder for the config's model type -- :func:`get_sezm_model` + for the DPA4/SeZM family (keeping its bridging/lora/compile/ + preset_out_bias rejections and ``exclude_types`` consistency check), + else :func:`get_standard_model` -- then re-classed through the + registered :class:`NativeSpinEnergyModel`. Eligibility is the atomic + model's own ``supports_native_spin()`` capability, not a descriptor-type + list -- so a bridging composition answers for itself. + + Parameters + ---------- + data : dict + The data to construct the model. Must carry a top-level ``"spin"`` + key with ``scheme == "native"``. + """ + data = copy.deepcopy(data) + spin_cfg = data.pop("spin") + data.setdefault("descriptor", {}) + # Expand index/symbol forms of ``use_spin`` against ``type_map`` into the + # per-type boolean list (pure; validates symbols). + use_spin = normalize_spin_use_spin(spin_cfg["use_spin"], data["type_map"]) + spin = Spin( + use_spin=use_spin, + virtual_scale=spin_cfg.get("virtual_scale", 1.0), + allow_missing_label=spin_cfg.get("allow_missing_label", False), + ) + data["descriptor"]["use_spin"] = use_spin + model_type = str(data.get("type", "standard")).lower() + backbone_builder = ( + get_sezm_model if model_type in ("dpa4", "sezm") else get_standard_model + ) + try: + backbone_model = backbone_builder(data) + except TypeError as err: + if "use_spin" not in str(err): + # Unrelated construction error (e.g. a bogus fitting kwarg): + # propagate with its real context instead of masking it as a + # capability failure. + raise + # A descriptor without native spin support rejects the injected + # ``use_spin`` keyword at construction; translate to the + # capability-gate error. + raise NotImplementedError( + "spin scheme 'native' requires a descriptor with native spin " + "support (supports_native_spin()); descriptor type " + f"{data['descriptor'].get('type')!r} does not accept `use_spin`" + ) from err + # The ATOMIC MODEL answers the capability -- it knows its own structure, + # so this holds for a plain descriptor+fitting model and for a bridging + # composition alike, with no assumption here about either. + if not backbone_model.atomic_model.supports_native_spin(): + raise NotImplementedError( + "spin scheme 'native' requires an atomic model declaring " + "supports_native_spin()" + ) + return NativeSpinEnergyModel(atomic_model_=backbone_model.atomic_model, spin=spin) def get_linear_model(model_params: dict) -> BaseModel: @@ -304,6 +414,11 @@ def get_model(data: dict) -> BaseModel: model_type = data.get("type", "standard") if model_type == "standard": if "spin" in data: + if str(data["spin"].get("scheme", "deepspin")) == "native": + # Descriptor-agnostic entry: any standard-typed config whose + # descriptor declares supports_native_spin() rides the + # native scheme with zero model/dispatch changes. + return get_native_spin_model(data) return get_spin_model(data) return get_standard_model(data) elif model_type == "linear_ener": diff --git a/deepmd/pt_expt/model/graph_lower.py b/deepmd/pt_expt/model/graph_lower.py index 788d9191a1..f196bd4fe5 100644 --- a/deepmd/pt_expt/model/graph_lower.py +++ b/deepmd/pt_expt/model/graph_lower.py @@ -40,11 +40,7 @@ def model_uses_graph_lower(model: Any) -> bool: except (AttributeError, NotImplementedError): return False - descriptor = getattr(getattr(model, "atomic_model", None), "descriptor", None) - uses_graph_lower = getattr(descriptor, "uses_graph_lower", None) - if uses_graph_lower is None: - return False - try: - return bool(uses_graph_lower()) - except (AttributeError, NotImplementedError): + atomic_model = getattr(model, "atomic_model", None) + if atomic_model is None: return False + return bool(atomic_model.uses_graph_lower()) diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index 3c77a62221..9d906c0874 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -16,9 +16,6 @@ from deepmd.dpmodel.atomic_model.base_atomic_model import ( BaseAtomicModel, ) -from deepmd.dpmodel.common import ( - get_xp_precision, -) from deepmd.dpmodel.model.make_model import make_model as make_model_dp from deepmd.dpmodel.output_def import ( OutputVariableDef, @@ -519,6 +516,7 @@ def forward_common_lower_graph( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None, + spin: torch.Tensor | None = None, comm_dict: dict | None = None, ) -> dict[str, torch.Tensor]: """Graph-native lower with autograd force/virial (dpa1/se_atten concat-tebd, attention included). @@ -583,8 +581,20 @@ def forward_common_lower_graph( multi-rank graphs the ghost rows are included; their values are inert under the owned-node mask). charge_spin - charge/spin conditioning. Ignored in PR-A; accepted for ABI - stability with charge/spin-conditioned descriptors. + Frame-level charge/spin FiLM conditioning, ``(nf, 2)`` or + ``None``, forwarded to the atomic model's + ``forward_common_atomic_graph`` (and, from there, the + descriptor's ``call_graph`` for descriptors that declare + ``supports_charge_spin``; currently DPA4 only). + spin + Per-node native spin, flat ``(N, 3)``, or ``None``. When given, + a SECOND autograd leaf is created next to ``edge_vec`` and + forwarded through the atomic-model chain to the descriptor + (native spin conditioning, e.g. DPA4/SeZM); the returned dict + additionally carries ``_derv_r_mag = -d_redu/dspin`` + for every ``r_differentiable`` reducible output. ``None`` + (default) is the existing, unconditioned graph lower with no + mag output. comm_dict MPI communication metadata for parallel inference. ``None`` (default) for non-parallel inference/training. Forwarded to @@ -612,6 +622,14 @@ def forward_common_lower_graph( # make edge_vec the autograd leaf for the energy backward edge_vec = edge_vec.detach().requires_grad_(True) + if spin is not None: + # second autograd leaf: force_mag = -dE/dspin (native spin). + # Deliberately a SEPARATE leaf/backward from edge_vec rather + # than a joint grad([edge_vec, spin]) call (as pt does in + # deepmd/pt/model/model/transform_output.py:288) -- this keeps + # edge_energy_deriv's signature untouched; see the second + # torch.autograd.grad call in fit_output_to_model_output_graph. + spin = spin.detach().requires_grad_(True) graph = NeighborGraph( n_node=n_node, edge_index=edge_index, @@ -626,7 +644,9 @@ def forward_common_lower_graph( ) # Level 2 emits force as a value through the inference-only custom # operator pipeline. Ineligible models use the autograd lower. - if not self.training and cuda_infer_level() >= 2: + # The fused pipeline has no mag output, so spin-conditioned models + # always take the autograd lower below. + if not self.training and cuda_infer_level() >= 2 and spin is None: fused = _fused_energy_force_graph(self, graph, atype, do_atomic_virial) if fused is not None: return fused @@ -636,6 +656,7 @@ def forward_common_lower_graph( fparam=fparam, aparam=aparam, charge_spin=charge_spin, + spin=spin, comm_dict=comm_dict, ) # ``forward_common_atomic_graph`` returns flat ``(N, *)`` output. @@ -647,13 +668,11 @@ def forward_common_lower_graph( do_atomic_virial=do_atomic_virial, create_graph=self.training, mask=atomic_ret["mask"] if "mask" in atomic_ret else None, - # Assemble force / virial in the descriptor compute precision - # (fp32 for an fp32 model) rather than the fp64 edge_vec leaf; - # the gradient content is only that precision, so the coarser - # scatter halves the atomic traffic at no accuracy cost. - force_precision=get_xp_precision( - torch, self.atomic_model.descriptor.precision - ), + spin_leaf=spin, + # Assemble force / virial in the INPUT (edge leaf) precision, + # consistent with the dpmodel backend's graph path and free + # of any assumption about the atomic model's internals. + force_precision=edge_vec.dtype, # Bound the per-node scatter by the INPUT node axis (the symbol # ``edge_index`` indexes into), not the re-derived fitting-output # shape -- avoids a CUDA out-of-bounds device-assert under @@ -699,10 +718,7 @@ def _resolve_graph_method( # for non-energy models (eager-only, output-agnostic). if "energy" not in self.atomic_output_def().keys(): return None - # Linear/ZBL atomic models have no single ``descriptor`` -> dense. - descriptor = getattr(self.atomic_model, "descriptor", None) - uses_graph_lower = getattr(descriptor, "uses_graph_lower", lambda: False) - if self.mixed_types() and uses_graph_lower(): + if self.mixed_types() and self.atomic_model.uses_graph_lower(): return "dense" return None @@ -715,6 +731,8 @@ def _call_common_graph( ap: torch.Tensor | None, method: str, do_atomic_virial: bool = False, + spin: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Carry-all graph forward with autograd force/virial (pt_expt override). @@ -739,6 +757,15 @@ def _call_common_graph( the carry-all builder, ``"dense"`` or ``"ase"``. do_atomic_virial whether to calculate the atomic virial. + spin + Per-local-atom native spin, ``(nf, nloc, 3)``, or ``None``. + Flattened to ``(N, 3)`` and forwarded into + :meth:`forward_common_lower_graph`, completing the seam + ``call_common`` (dpmodel, shared) opens for the graph route. + charge_spin + Frame-level charge/spin FiLM conditioning, ``(nf, 2)`` or + ``None``. Unflattened (per-frame) and forwarded unchanged into + :meth:`forward_common_lower_graph`. Returns ------- @@ -752,18 +779,20 @@ def _call_common_graph( # check only protects the default (None) path; an EXPLICIT # neighbor_graph_method would otherwise reach the builders for # descriptors without a graph lower. - descriptor = getattr(self.atomic_model, "descriptor", None) - uses_graph_lower = getattr(descriptor, "uses_graph_lower", lambda: False) - if not (self.mixed_types() and uses_graph_lower()): + if not (self.mixed_types() and self.atomic_model.uses_graph_lower()): raise NotImplementedError( "neighbor_graph_method requires a mixed_types descriptor with a " "graph lower (e.g. dpa1 attn_layer=0)" ) rcut = self.get_rcut() + # CSR pre-sort serves the compressed-DPA1 fused kernels only; + # probe via the atomic model's own descriptor when it has one + # (compositions have none and never take the fused path). + _desc = getattr(self.atomic_model, "descriptor", None) with_csr = ( not self.training and cuda_infer_level() >= 1 - and bool(getattr(descriptor, "geo_compress", False)) + and bool(getattr(_desc, "geo_compress", False)) ) pair_excl = getattr(self.atomic_model, "pair_excl", None) ng = _build_graph_for_method( @@ -771,8 +800,9 @@ def _call_common_graph( ) nf, nloc = atype.shape[:2] atype_flat = atype.reshape(nf * nloc) - # graph-lower ABI: aparam is FLAT on the node axis, (N, nda). + # graph-lower ABI: aparam/spin are FLAT on the node axis, (N, nda)/(N, 3). ap_flat = ap.reshape(nf * nloc, ap.shape[-1]) if ap is not None else None + spin_flat = spin.reshape(nf * nloc, 3) if spin is not None else None model_predict = self.forward_common_lower_graph( atype_flat, ng.n_node, @@ -788,6 +818,8 @@ def _call_common_graph( do_atomic_virial=do_atomic_virial, fparam=fp, aparam=ap_flat, + spin=spin_flat, + charge_spin=charge_spin, ) # ``forward_common_lower_graph`` returns flat ``(N, *)`` per-atom # outputs (N = nf * nloc for a carry-all rectangular graph). @@ -987,6 +1019,7 @@ def forward_common_lower_graph_exportable( aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None, + spin: torch.Tensor | None = None, destination_sorted: bool = False, **make_fx_kwargs: Any, ) -> torch.nn.Module: @@ -1019,6 +1052,15 @@ def forward_common_lower_graph_exportable( destination-major and ``destination_order`` is identity. fparam, aparam, do_atomic_virial, charge_spin As in ``forward_common_lower_graph``. + spin + Per-node native spin, flat ``(N, 3)``, or ``None``. Threaded + through to ``forward_common_lower_graph`` the same way as + ``charge_spin``; when given, the trace additionally carries a + SECOND autograd leaf so the returned dict carries + ``_derv_r_mag`` for every ``r_differentiable`` reducible + output. ``None`` (default) is the existing, unconditioned + trace with no mag output -- used by every non-spin caller of + this generic (descriptor-agnostic) exportable. **make_fx_kwargs Extra keyword arguments forwarded to ``make_fx`` (e.g. ``tracing_mode="symbolic"``). @@ -1029,8 +1071,8 @@ def forward_common_lower_graph_exportable( A traced module whose ``forward`` accepts ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, - source_row_ptr, fparam, aparam, charge_spin)`` and returns a - dict with the same internal keys as + source_row_ptr, fparam, aparam, charge_spin, spin)`` and + returns a dict with the same internal keys as ``forward_common_lower_graph``. """ validate_graph_csr_for_export( @@ -1045,7 +1087,70 @@ def forward_common_lower_graph_exportable( ) model = self - def fn( + # ``spin`` is a traced INPUT only when a real spin tensor is + # given (native-spin models). For every non-spin caller + # (``spin is None``) the traced ``fn`` must have the SAME arity as + # before native spin existed -- otherwise the outer energy trace, + # which threads ``charge_spin`` but no ``spin``, would call this + # traced module missing a ``spin`` argument. So the ``spin=None`` + # branch traces the original 13-input closure (spin captured as a + # ``None`` constant), and the spin branch traces a 14-input + # closure with ``spin`` at the tail. + if spin is None: + + def fn( + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + # forward_common_lower_graph creates the autograd leaf from + # edge_vec internally, so no outer detach/requires_grad_ + # here (it would only add spurious ops to the traced graph). + return model.forward_common_lower_graph( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + destination_sorted=destination_sorted, + do_atomic_virial=do_atomic_virial, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + ) + + return make_fx(fn, **make_fx_kwargs)( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + ) + + def fn_spin( atype: torch.Tensor, n_node: torch.Tensor, n_local: torch.Tensor, @@ -1059,10 +1164,12 @@ def fn( fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, + spin: torch.Tensor | None, ) -> dict[str, torch.Tensor]: - # forward_common_lower_graph creates the autograd leaf from - # edge_vec internally, so no outer detach/requires_grad_ here - # (it would only add spurious ops to the traced graph). + # forward_common_lower_graph creates the autograd leaf(s) from + # edge_vec AND spin internally, so no outer + # detach/requires_grad_ here (it would only add spurious ops + # to the traced graph). return model.forward_common_lower_graph( atype, n_node, @@ -1079,9 +1186,10 @@ def fn( fparam=fparam, aparam=aparam, charge_spin=charge_spin, + spin=spin, ) - return make_fx(fn, **make_fx_kwargs)( + return make_fx(fn_spin, **make_fx_kwargs)( atype, n_node, n_local, @@ -1095,6 +1203,7 @@ def fn( fparam, aparam, charge_spin, + spin, ) def forward_common_lower_exportable_with_comm( diff --git a/deepmd/pt_expt/model/model.py b/deepmd/pt_expt/model/model.py index a4b18d3d41..a16ae9b31b 100644 --- a/deepmd/pt_expt/model/model.py +++ b/deepmd/pt_expt/model/model.py @@ -1,116 +1,21 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Any, -) - -from deepmd.dpmodel.model.base_model import ( - make_base_model, -) -from deepmd.utils.version import ( - check_version_compatibility, -) - - -class BaseModel(make_base_model()): - """Base class for pt_expt models. - - Provides the plugin registry so that model classes can be - registered with ``@BaseModel.register("ener")`` etc. - - See Also - -------- - deepmd.dpmodel.model.base_model.BaseBaseModel - Backend-independent BaseModel class. - """ - - # The pt backend's ``SeZMModel`` (model_type "SeZM", aliases dpa4/sezm) - # serialises with a *model-level wrapper*: ``{type: "SeZM", - # atomic_model: , bridging_method, bridging_r_*, lora}``, - # and its atomic model uses ``type: "sezm_atomic"`` carrying pt-only - # extras (``dens_fitting``/``active_mode`` plus a ``dens_force_rmsd`` - # @variable). pt_expt builds the equivalent DPA4 model via the generic - # ``make_model`` path, whose ``serialize()`` emits the standard atomic - # dict directly (``type: "standard"``). To load a pt-trained checkpoint - # into pt_expt (the serialization-compat / checkpoint-interop - # requirement), recognise the wrapper, reject the pt-only features pt_expt - # does not implement (when they are non-default), strip the rest, and - # delegate to the standard path. The nested descriptor/fitting dicts are - # already backend-agnostic dpmodel serializations and pass through intact. - _SEZM_MODEL_TYPES = frozenset({"sezm", "dpa4"}) - _SEZM_ATOMIC_TYPES = frozenset({"sezm_atomic"}) - - @classmethod - def deserialize(cls, data: dict[str, Any]) -> "BaseModel": - model_type = str(data.get("type", "standard")) - if model_type.lower() in cls._SEZM_MODEL_TYPES: - return cls.deserialize(cls._unwrap_pt_sezm_model(data)) - if model_type.lower() in cls._SEZM_ATOMIC_TYPES: - return cls.deserialize(cls._normalize_pt_sezm_atomic(data)) - return super().deserialize(data) - - @staticmethod - def _unwrap_pt_sezm_model(data: dict[str, Any]) -> dict[str, Any]: - """Unwrap pt's ``SeZMModel`` serialization to the inner atomic dict.""" - data = data.copy() - # The pt SeZM model wrapper serialises with ``@version`` 1. Validate - # before discarding it so a future incompatible wrapper schema is not - # silently mis-deserialized (the wrapper only carries the guarded - # bridging/lora extras below, so the accepted range is narrow). - check_version_compatibility(int(data.get("@version", 1)), 1, 1) - bridging_method = str(data.get("bridging_method", "none")).lower() - if bridging_method not in ("none", ""): - raise NotImplementedError( - "Deserializing a pt SeZM/DPA4 checkpoint with " - f"`bridging_method`={data.get('bridging_method')!r} is not " - "supported in pt_expt." - ) - if data.get("lora") is not None: - raise NotImplementedError( - "Deserializing a pt SeZM/DPA4 checkpoint with `lora` is " - "not supported in pt_expt." - ) - atomic_model = data.get("atomic_model") - if atomic_model is None: - raise ValueError( - "SeZM/DPA4 model data is missing the 'atomic_model' entry." - ) - return atomic_model - - @staticmethod - def _normalize_pt_sezm_atomic(data: dict[str, Any]) -> dict[str, Any]: - """Convert a pt ``sezm_atomic`` dict to a standard atomic dict. - - Strips the pt-only ``dens`` head state (``dens_fitting`` / - ``active_mode`` / the ``dens_force_rmsd`` @variable) and rewrites the - ``type``/``@version`` so the generic dpmodel atomic-model deserialize - accepts it. A non-energy active mode or a populated dens head is - rejected because pt_expt only implements the energy path. - """ - data = data.copy() - # pt emits ``@version`` 3 for ``sezm_atomic``; the standard dpmodel - # atomic-model deserialize requires exactly 2. The only schema delta - # between the two is the stripped ``dens`` state below, so coercion is - # safe for the known-compatible range {2, 3}. Validate the incoming - # version BEFORE coercing so a future incompatible pt schema (e.g. - # ``@version`` 4) is rejected loudly instead of mis-deserialized. - check_version_compatibility(int(data.get("@version", 2)), 3, 2) - if data.pop("dens_fitting", None) is not None: - raise NotImplementedError( - "Deserializing a pt SeZM/DPA4 checkpoint with a `dens` " - "fitting head is not supported in pt_expt." - ) - active_mode = data.pop("active_mode", None) - if active_mode not in (None, "ener"): - raise NotImplementedError( - f"Deserializing a pt SeZM/DPA4 checkpoint in active_mode " - f"{active_mode!r} is not supported in pt_expt (energy only)." - ) - variables = data.get("@variables") - if isinstance(variables, dict): - data["@variables"] = { - k: v for k, v in variables.items() if k in ("out_bias", "out_std") - } - # The standard dpmodel atomic-model deserialize checks @version == 2. - data["@version"] = 2 - data["type"] = "standard" - return data +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.model.base_model import ( + make_base_model, +) + + +class BaseModel(make_base_model()): + """Base class for pt_expt models. + + Provides the plugin registry so that model classes can be registered + with ``@BaseModel.register("ener")`` etc. Deserialization is pure + registry dispatch: descriptor-family-specific wire formats (e.g. the + pt SeZM checkpoint layouts) are owned by the family's registered model + class (see ``deepmd.pt_expt.model.dpa4_model.DPA4EnergyModel``), never + by this base. + + See Also + -------- + deepmd.dpmodel.model.base_model.BaseBaseModel + Backend-independent BaseModel class. + """ diff --git a/deepmd/pt_expt/model/native_spin_model.py b/deepmd/pt_expt/model/native_spin_model.py new file mode 100644 index 0000000000..0d63e7e530 --- /dev/null +++ b/deepmd/pt_expt/model/native_spin_model.py @@ -0,0 +1,440 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""pt_expt native-spin energy model (NeighborGraph route, autograd force_mag).""" + +from typing import ( + Any, +) + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.dpmodel.model.native_spin_model import ( + make_native_spin_model, +) +from deepmd.pt_expt.model.ener_model import ( + EnergyModel, +) +from deepmd.pt_expt.model.model import ( + BaseModel, +) + + +@BaseModel.register("native_spin") +class NativeSpinEnergyModel(make_native_spin_model(EnergyModel)): + """pt_expt native-spin energy model. + + ``make_native_spin_model`` applied to THIS backend's + :class:`~deepmd.pt_expt.model.ener_model.EnergyModel` (construction, + output defs, serialization all come from the factory; the deserialize + closure rebuilds through the pt_expt model class, so a registry round + trip yields a real ``torch.nn.Module``), plus two torch-specific + overrides: + + - :meth:`forward`: the pt_expt ``call_common`` produces REAL autograd + ``energy_derv_r``/``energy_derv_r_mag``/``energy_derv_c_redu`` + tensors, unlike the dpmodel factory's energy-only ``call`` (which is + restricted to ``force``/``force_mag``/``virial`` as ``None`` + placeholders because dpmodel has no autograd). + - :meth:`forward_lower_graph_exportable`: the graph-spin ``.pt2`` + positional ABI (``spin`` at index 10) over the inherited + ``forward_common_lower_graph_exportable``. + """ + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + spin: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Return native-spin model predictions with public output keys. + + Parameters + ---------- + coord + The coordinates of the atoms. shape: nf x (nloc x 3) + atype + The type of atoms. shape: nf x nloc + spin + The per-local-atom spin. shape: nf x (nloc x 3) + box + The simulation box. shape: nf x 9 + fparam + frame parameter. nf x ndf + aparam + atomic parameter. nf x nloc x nda + do_atomic_virial + If calculate the atomic virial. + charge_spin + Frame-level charge/spin conditioning, shape nf x 2. Accepted for + call-signature compatibility with ``ModelWrapper.forward`` (which + always forwards this keyword); charge-spin FiLM combined with + native spin is rejected at construction time + (``add_chg_spin_ebd`` on the descriptor), so this is always + ``None`` in practice for a model this class can build. + + Returns + ------- + ret_dict + The result dict with keys ``atom_energy``, ``energy``, + ``force``, ``force_mag``, ``virial``, ``mask_mag``, and + (when ``do_atomic_virial``) ``atom_virial``. ``force`` and + ``force_mag`` are real autograd tensors (``-dE/dcoord`` and + ``-dE/dspin``), NOT placeholders. + """ + # Default neighbor_graph_method: pt_expt's default-flip picks the fast + # graph builder (dpmodel's call forces the O(N^2) dense one). + model_ret = self.call_common( + coord, + atype, + box=box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, + spin=spin, + ) + # Same shared translation as dpmodel's call, on real autograd tensors. + return self._translate_eager_call( + model_ret, atype, do_atomic_virial=do_atomic_virial + ) + + def forward_lower_graph_exportable( + self, + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + spin: torch.Tensor, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + destination_sorted: bool = False, + **make_fx_kwargs: Any, + ) -> torch.nn.Module: + """Trace the graph-spin lower into an exportable module. + + THIS METHOD OWNS the positional ``.pt2`` ABI for graph-spin models + (mirrored verbatim by the C++ / serialization seams): ``spin`` sits + at index 10, right after the shared ``NeighborGraph`` CSR block and + before the conditional ``fparam``/``aparam`` tail; the conditional + ``charge_spin`` tail follows at index 13 (combined native-spin + + charge-spin FiLM models, review 3638047227; ``None`` otherwise). + :meth:`forward_lower_graph_exportable_with_comm` extends this SAME + prefix with the 8 comm tensors for multi-rank. + + Two-layer make_fx trace, mirroring + :meth:`~deepmd.pt_expt.model.ener_model.EnergyModel.forward_lower_graph_exportable`: + the inner layer + (the inherited + :meth:`~deepmd.pt_expt.model.make_model.make_model.forward_common_lower_graph_exportable`) + traces ``forward_common_lower_graph`` + with ``spin`` as a SECOND autograd leaf next to ``edge_vec`` + (``charge_spin`` is a frame-level conditioning input, not a leaf); + this outer layer re-traces with the PUBLIC positional ABI above and + translates the internal fitting keys to the public output keys via + the shared :meth:`_translate_eager_call` (the same single owner as + the eager :meth:`forward`, so ``mask_mag`` is emitted here too). + + Parameters + ---------- + atype + (N,) flat local-plus-halo atom types, ``N == sum(n_node)``. + n_node + (nf,) per-frame total node counts. + n_local + (nf,) per-frame owned node counts. + edge_index + (2, E) ``[src, dst]`` edge endpoints (flat local indices). + edge_vec + (E, 3) neighbor-minus-center edge vectors (sample for tracing). + edge_mask + (E,) valid-edge mask (sample for tracing). + destination_order + (E,) destination-grouped edge permutation. + destination_row_ptr + (N + 1,) destination CSR offsets. + source_order + (E,) source-grouped edge permutation. + source_row_ptr + (N + 1,) source CSR offsets. + spin + (N, 3) per-node native spin (sample for tracing). ALWAYS present + in this ABI, unlike the energy model's optional ``charge_spin``. + fparam + Frame parameter, ``(nf, ndf)``, or ``None`` when + ``dim_fparam == 0``. + aparam + Atomic parameter, ``(N, nda)``, or ``None`` when + ``dim_aparam == 0``. + charge_spin + Frame-level charge/spin FiLM conditioning, ``(nf, + dim_chg_spin)``, or ``None`` when the descriptor has no + ``add_chg_spin_ebd`` (conditional tail, slot 13). + do_atomic_virial + Whether to also return ``atom_virial``. + destination_sorted + Static export-time assertion that the payload is + destination-major and ``destination_order`` is identity. + **make_fx_kwargs + Extra keyword arguments forwarded to ``make_fx`` (e.g. + ``tracing_mode="symbolic"``). + + Returns + ------- + torch.nn.Module + A traced module whose ``forward`` accepts ``(atype, n_node, + n_local, edge_index, edge_vec, edge_mask, destination_order, + destination_row_ptr, source_order, source_row_ptr, spin, + fparam, aparam, charge_spin)`` and returns a dict with the + public keys ``atom_energy``, ``energy``, ``force``, + ``force_mag``, ``virial``, ``mask_mag``, and (when + ``do_atomic_virial``) ``atom_virial``. + """ + traced = self.forward_common_lower_graph_exportable( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, + spin=spin, + destination_sorted=destination_sorted, + **make_fx_kwargs, + ) + + def fn( + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + spin: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + model_ret = traced( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + spin, + ) + # Same single-owner translation as the eager forward, so the + # exported .pt2 emits mask_mag too and consumers read it. + return self._translate_eager_call( + model_ret, atype, do_atomic_virial=do_atomic_virial + ) + + return make_fx(fn, **make_fx_kwargs)( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + spin, + fparam, + aparam, + charge_spin, + ) + + def forward_lower_graph_exportable_with_comm( + self, + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + spin: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + send_list: torch.Tensor, + send_proc: torch.Tensor, + recv_proc: torch.Tensor, + send_num: torch.Tensor, + recv_num: torch.Tensor, + communicator: torch.Tensor, + nlocal: torch.Tensor, + nghost: torch.Tensor, + do_atomic_virial: bool = False, + **make_fx_kwargs: Any, + ) -> torch.nn.Module: + """Trace the multi-rank graph-spin lower into an exportable module. + + The with-comm counterpart of + :meth:`forward_lower_graph_exportable`: same positional prefix, + ``spin`` still at index 10, then the 8 comm tensors appended after + the conditional ``fparam``/``aparam``/``charge_spin`` tail (indices + 14-21) -- exactly how + :meth:`~deepmd.pt_expt.model.ener_model.EnergyModel.forward_lower_graph_exportable_with_comm` + appends them for the energy model. + + ``spin`` is the EXTENDED per-node spin ``(N, 3)``: ghost rows carry + their owner's spin, delivered by the LAMMPS ``sp`` forward-comm + before the call, so the descriptor's spin embedding sees the same + value on every rank that holds the node. The per-block ghost + FEATURE refresh rides ``deepmd_export::border_op`` exactly as in the + energy model; spin needs no border exchange of its own because it is + an input, not a derived feature. + + Single make_fx trace (the energy with-comm precedent), unlike the + two-layer trace of the non-comm spin path: the comm-dict packing, + the ``forward_common_lower_graph`` call with ``spin`` as a second + autograd leaf, and the public-key translation all live in one traced + ``fn``. + + Parameters + ---------- + atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, source_row_ptr, spin, fparam, aparam, charge_spin + As in :meth:`forward_lower_graph_exportable`. + send_list, send_proc, recv_proc, send_num, recv_num, communicator, nlocal, nghost + The 8 comm tensors, packed into ``comm_dict`` inside the traced + function. Same runtime device contract as the energy model's: + ALL 8 stay on CPU (host control metadata for ``border_op``), + while the device-side owned count is the separate ``n_local`` + input at slot 2. + do_atomic_virial + Whether to also return ``atom_virial``. + **make_fx_kwargs + Extra keyword arguments forwarded to ``make_fx``. + + Returns + ------- + torch.nn.Module + A traced module accepting the 22-input ABI above and returning + the same public keys as :meth:`forward_lower_graph_exportable` + (``atom_energy``, ``energy``, ``force``, ``force_mag``, + ``virial``, ``mask_mag``, plus ``atom_virial`` when requested). + """ + model = self + + def fn( + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + spin: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + send_list: torch.Tensor, + send_proc: torch.Tensor, + recv_proc: torch.Tensor, + send_num: torch.Tensor, + recv_num: torch.Tensor, + communicator: torch.Tensor, + nlocal: torch.Tensor, + nghost: torch.Tensor, + ) -> dict[str, torch.Tensor]: + comm_dict = { + "send_list": send_list, + "send_proc": send_proc, + "recv_proc": recv_proc, + "send_num": send_num, + "recv_num": recv_num, + "communicator": communicator, + "nlocal": nlocal, + "nghost": nghost, + } + model_ret = model.forward_common_lower_graph( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + destination_sorted=True, + do_atomic_virial=do_atomic_virial, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + spin=spin, + comm_dict=comm_dict, + ) + # Same single-owner translation as the eager forward and the + # non-comm lower, so mask_mag is emitted here too. + return self._translate_eager_call( + model_ret, atype, do_atomic_virial=do_atomic_virial + ) + + return make_fx(fn, **make_fx_kwargs)( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + spin, + fparam, + aparam, + charge_spin, + send_list, + send_proc, + recv_proc, + send_num, + recv_num, + communicator, + nlocal, + nghost, + ) diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 7db988364b..697bf46dfb 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -261,6 +261,22 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: "aparam", _model.get_dim_aparam(), atomic=True, must=True ) ) + if _model.has_spin(): + # ``model.spin.allow_missing_label`` relaxes the spin label from + # mandatory to optional with a zero default, so a system without a + # ``spin`` file is filled with zeros rather than rejected. Mirrors + # ``deepmd.pt.train.training.get_additional_data_requirement``. + # Every spin model wrapper carries a ``spin`` attribute. + allow_missing_spin = _model.spin.allow_missing_label + additional_data_requirement.append( + DataRequirementItem( + "spin", + ndof=3, + atomic=True, + must=not allow_missing_spin, + default=0.0, + ) + ) if _model.has_chg_spin_ebd(): has_default_cs = _model.has_default_chg_spin() if has_default_cs: @@ -1852,10 +1868,9 @@ def _raise_if_full_validation_unsupported( "training; multi-task training is not supported." ) - has_spin = getattr(self.models[DEFAULT_TASK_KEY], "has_spin", False) - if callable(has_spin): - has_spin = has_spin() - if has_spin or isinstance(self.loss, EnergySpinLoss): + if self.models[DEFAULT_TASK_KEY].has_spin() or isinstance( + self.loss, EnergySpinLoss + ): raise ValueError( "validating.full_validation only supports single-task energy " "training; spin-energy training is not supported." diff --git a/deepmd/pt_expt/train/wrapper.py b/deepmd/pt_expt/train/wrapper.py index 1509924e93..f59b707217 100644 --- a/deepmd/pt_expt/train/wrapper.py +++ b/deepmd/pt_expt/train/wrapper.py @@ -139,6 +139,7 @@ def forward( self, coord: torch.Tensor, atype: torch.Tensor, + spin: torch.Tensor | None = None, box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, @@ -164,6 +165,13 @@ def forward( "aparam": aparam, "charge_spin": charge_spin, } + # ``spin`` (native or virtual-atom magnetic moment) is only accepted + # by spin-capable model forward()s; mirrors + # ``deepmd.pt.train.wrapper.ModelWrapper.forward``'s ``has_spin`` gate + # so non-spin models (whose forward() has no ``spin`` parameter) are + # never called with an unexpected keyword argument. + if self.model[task_key].has_spin(): + input_dict["spin"] = spin if self.inference_only: with self._frozen_parameter_context(): diff --git a/deepmd/pt_expt/utils/__init__.py b/deepmd/pt_expt/utils/__init__.py index 4e637e5d4f..d30f288909 100644 --- a/deepmd/pt_expt/utils/__init__.py +++ b/deepmd/pt_expt/utils/__init__.py @@ -11,6 +11,9 @@ AtomExcludeMask, PairExcludeMask, ) +from .inter_potential import ( + InterPotential, +) from .network import ( NetworkCollection, ) @@ -35,6 +38,7 @@ __all__ = [ "AtomExcludeMask", + "InterPotential", "NetworkCollection", "PairExcludeMask", "TypeEmbedNet", diff --git a/deepmd/pt_expt/utils/edge_schema.py b/deepmd/pt_expt/utils/edge_schema.py index baa7eca2dd..4e916abea0 100644 --- a/deepmd/pt_expt/utils/edge_schema.py +++ b/deepmd/pt_expt/utils/edge_schema.py @@ -1,5 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Edge-vector neighbor-list helpers for SeZM-style models.""" +"""Edge-vector neighbor-list helpers for SeZM-style models. + +LEGACY: serves only the deprecated pt-backend edge_vec .pt2 schema; +scheduled for removal with that rail. +""" from __future__ import ( annotations, diff --git a/deepmd/pt_expt/utils/inter_potential.py b/deepmd/pt_expt/utils/inter_potential.py new file mode 100644 index 0000000000..0647e9a73f --- /dev/null +++ b/deepmd/pt_expt/utils/inter_potential.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""pt_expt wrapper for the analytical bridging pair potential.""" + +from typing import ( + Any, +) + +from deepmd.dpmodel.atomic_model.inter_potential import ( + InterPotential as InterPotentialDP, +) +from deepmd.pt_expt.common import ( + register_dpmodel_mapping, + torch_module, +) + + +@torch_module +class InterPotential(InterPotentialDP): + def forward(self, *args: Any, **kwargs: Any) -> Any: + return self.call(*args, **kwargs) + + +# InterPotential carries no trainable state (only the constant per-type +# atomic-number table, derived from the constructor arguments), so it +# implements no serialize()/deserialize(); rebuild it fresh from +# (type_map, mode). +register_dpmodel_mapping( + InterPotentialDP, + lambda v: InterPotential(type_map=v.type_map, mode=v.mode), +) diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 84d9c53831..897d91cec3 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -121,7 +121,9 @@ def _metadata_value_to_json(value: Any) -> Any: return value -def _needs_with_comm_artifact(model: torch.nn.Module) -> bool: +def _needs_with_comm_artifact( + model: torch.nn.Module, lower_kind: str = "nlist" +) -> bool: """Return ``True`` if the model needs a "with-comm" AOTI artifact compiled. The with-comm artifact carries the per-layer ``deepmd_export::border_op`` @@ -136,15 +138,78 @@ def _needs_with_comm_artifact(model: torch.nn.Module) -> bool: descriptor classes implement explicitly. Returns ``False`` defensively when the model has no single descriptor (linear/zbl/frozen) or when the method is somehow missing or raises. + + Not every lower path that needs cross-rank exchange implements it: DPA4's + graph lower carries a real per-layer ``border_op`` exchange, but its + dense (nlist) lower's adapter raises on ``comm_dict``. ``lower_kind`` + selects which lower is being traced so the gate can consult the + per-lower capability instead of assuming both lowers agree. Non-graph + kinds additionally check ``descriptor.dense_lower_supports_comm()`` + (absent on descriptors, such as dpa2/dpa3, whose dense lower always + supports comm — treated as ``True``). + + Native spin participates on the GRAPH lower, matching pt's + ``SeZMModel.supports_edge_parallel`` (which ``SeZMNativeSpinModel`` does + not override): the spin input is per-node and its ghost rows arrive via + the LAMMPS ``sp`` forward-comm, so nothing about spin needs its own + cross-rank exchange -- the per-block ghost FEATURE refresh is the same + ``border_op`` the energy model uses. It is excluded only on the dense + (nlist) lower, which has no spin with-comm wrapper. + + Parameters + ---------- + model : torch.nn.Module + The wrapped pt_expt model. + lower_kind : str + Which lower is being traced/frozen: ``"graph"`` or ``"nlist"`` + (dense). Defaults to ``"nlist"``. + + Returns + ------- + bool + Whether a with-comm artifact should be built for this lower kind. """ - desc = getattr(getattr(model, "atomic_model", None), "descriptor", None) - if desc is None or not hasattr(desc, "has_message_passing_across_ranks"): + from deepmd.dpmodel.model.native_spin_model import ( + NativeSpinModelKind, + ) + + # Cross-backend family test: the dpmodel and pt_expt concrete classes + # are parallel factory products with no subclass relation, so the shared + # marker base -- not a concrete class -- is the membership check. + # Native spin rides the GRAPH lower only; its dense lower has no + # with-comm wrapper at all. + if isinstance(model, NativeSpinModelKind) and lower_kind != "graph": return False - try: - return bool(desc.has_message_passing_across_ranks()) - except (AttributeError, NotImplementedError): + + # Analytical bridging models are single-rank only (pt's + # ``supports_edge_parallel() == False`` contract: ZBL + SFPG fold each + # node's full outgoing-edge set, which a single rank cannot observe for + # ghost owners) -- never compile a with-comm artifact for them. + from deepmd.dpmodel.atomic_model.linear_atomic_model import ( + LinearEnergyAtomicModel, + ) + + atomic_model = getattr(model, "atomic_model", None) + if isinstance(atomic_model, LinearEnergyAtomicModel): + # Compositions (e.g. analytical bridging: learned + InterPotential) + # are single-rank on the graph route: per-edge analytical terms fold + # each node's full edge set, which a single rank cannot observe for + # ghost owners (pt's supports_edge_parallel()==False rationale). return False + desc = getattr(getattr(model, "atomic_model", None), "descriptor", None) + if desc is None or not desc.has_message_passing_across_ranks(): + return False + if lower_kind == "graph": + return True + # Non-graph kinds trace the DENSE with-comm wrapper; a descriptor whose + # dense lower has no comm implementation (DPA4: the dense adapter raises + # on comm_dict) must not emit a dead or untraceable dense artifact. + # Descriptors without the method (dpa2/dpa3/...) implement dense comm — + # it is their production multi-rank path. + dense_ok = getattr(desc, "dense_lower_supports_comm", None) + return True if dense_ok is None else bool(dense_ok()) + def check_graph_trace_torch_version(model: torch.nn.Module) -> None: """Fail fast when the graph trace needs unbacked-SymInt support torch lacks. @@ -181,8 +246,7 @@ def check_graph_trace_torch_version(model: torch.nn.Module) -> None: running torch is older than 2.6. """ desc = getattr(getattr(model, "atomic_model", None), "descriptor", None) - uses_pairs = getattr(desc, "uses_compact_edge_pairs", None) - if uses_pairs is None or not uses_pairs(): + if desc is None or not desc.uses_compact_edge_pairs(): return version = torch.__version__.split("+")[0] major_minor = tuple(int(p) for p in version.split(".")[:2] if p.isdigit()) @@ -383,6 +447,7 @@ def build_synthetic_graph_inputs( want_fparam: bool = True, want_aparam: bool = True, want_charge_spin: bool = True, + want_spin: bool = False, ) -> tuple[torch.Tensor | None, ...]: """Build a synthetic carry-all ``NeighborGraph`` for graph-lower tracing. @@ -399,7 +464,13 @@ def build_synthetic_graph_inputs( ``forward_(common_)lower_graph``: ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, source_row_ptr, fparam, aparam, - charge_spin)``. + charge_spin)`` -- or, when ``want_spin=True``, the native-spin ABI + (:meth:`~deepmd.pt_expt.model.native_spin_model.NativeSpinEnergyModel.forward_lower_graph_exportable`): + ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, + destination_row_ptr, source_order, source_row_ptr, spin, fparam, aparam)`` + -- ``spin`` replaces ``charge_spin`` at the tail AND moves to slot 10 + (before ``fparam``/``aparam``); there is no ``charge_spin`` slot at all + (native spin rejects ``add_chg_spin_ebd`` at build). The system (``rng(42)``, ``box = rcut*3``, centered coords, ``atype[:, i] = i % ntypes``) is identical for both callers; the only two former differences @@ -434,6 +505,14 @@ def build_synthetic_graph_inputs( Whether to emit the optional conditioning tensor when its ``dim > 0``. Export passes the defaults (``True`` = include if present); training passes ``x is not None`` so the traced branch matches the run-time call. + want_spin : bool + Build the native-spin ABI instead of the regular energy ABI: insert + a ``(N, 3)`` sample spin tensor at slot 10 and drop the + ``charge_spin`` slot. The sample is a small NON-ZERO deterministic + value (``0.1 + 0.05 * arange(...)``, NOT ``torch.zeros`` -- an + all-zero spin leaf can hit degenerate branches, e.g. a + ``norm(spin) == 0`` special case, in the equivariant spin + embedding). """ import deepmd.pt_expt.utils.env as _env from deepmd.dpmodel.utils.neighbor_graph import ( @@ -487,15 +566,46 @@ def build_synthetic_graph_inputs( if (want_aparam and dim_aparam > 0) else None ) + # Keep total and owned counts value-distinct during tracing so export does + # not specialize the multi-rank ownership relation to ``n_local == n_node``. + n_local = torch.clamp(graph.n_node - 1, min=1) + + if want_spin: + # Native-spin ABI: spin at slot 10 (before fparam/aparam), then the + # conditional charge_spin tail at slot 13 (combined native-spin + + # charge-spin FiLM models). A small NON-ZERO deterministic spin + # sample -- NOT torch.zeros (see the docstring's want_spin entry). + n_node_total = nframes * nloc + spin = ( + 0.1 + 0.05 * torch.arange(n_node_total * 3, dtype=dtype, device=device) + ).reshape(n_node_total, 3) + charge_spin = ( + torch.zeros(nframes, dim_chg_spin, dtype=dtype, device=device) + if (want_charge_spin and dim_chg_spin > 0) + else None + ) + return ( + atype_t.reshape(-1), + graph.n_node, + n_local, + graph.edge_index, + graph.edge_vec.to(edge_dtype), + graph.edge_mask, + graph.destination_order, + graph.destination_row_ptr, + graph.source_order, + graph.source_row_ptr, + spin, + fparam, + aparam, + charge_spin, + ) + charge_spin = ( torch.zeros(nframes, dim_chg_spin, dtype=dtype, device=device) if (want_charge_spin and dim_chg_spin > 0) else None ) - # Keep total and owned counts value-distinct during tracing so export does - # not specialize the multi-rank ownership relation to ``n_local == n_node``. - n_local = torch.clamp(graph.n_node - 1, min=1) - return ( atype_t.reshape(-1), graph.n_node, @@ -649,6 +759,7 @@ def count_synthetic_graph_edges( def _build_graph_dynamic_shapes( *sample_inputs: torch.Tensor | None, + is_native_spin: bool = False, ) -> tuple: """Build dynamic-shape specifications for the graph-form forward_lower export. @@ -663,18 +774,24 @@ def _build_graph_dynamic_shapes( Parameters ---------- *sample_inputs : torch.Tensor | None - ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, - destination_order, destination_row_ptr, source_order, source_row_ptr, - fparam, aparam, charge_spin)`` — 13 entries matching - ``forward_lower_graph_exportable``. + Regular (energy) ABI: ``(atype, n_node, n_local, edge_index, + edge_vec, edge_mask, destination_order, destination_row_ptr, + source_order, source_row_ptr, fparam, aparam, charge_spin)`` — 13 + entries matching ``forward_lower_graph_exportable``. Native-spin ABI + (``is_native_spin=True``): same shared CSR block (slots 0-9), but + slot 10 is ``spin`` (mandatory, node-axis-shaped), slot 11 + ``fparam``, slot 12 ``aparam``, slot 13 the conditional + ``charge_spin`` tail (see + ``NativeSpinEnergyModel.forward_lower_graph_exportable``). + is_native_spin : bool + Whether ``sample_inputs`` follows the native-spin positional ABI + (spin at slot 10) instead of the regular energy ABI (charge_spin at + slot 12). """ - fparam = sample_inputs[10] - aparam = sample_inputs[11] - charge_spin = sample_inputs[12] nframes_dim = torch.export.Dim("nframes", min=1) n_node_total_dim = torch.export.Dim("n_node_total", min=1) nedge_dim = torch.export.Dim("nedge", min=2) - return ( + base = ( {0: n_node_total_dim}, # atype: (N,) {0: nframes_dim}, # n_node: (nf,) {0: nframes_dim}, # n_local: (nf,) @@ -685,6 +802,26 @@ def _build_graph_dynamic_shapes( {0: n_node_total_dim + 1}, # destination_row_ptr: (N + 1,) {0: nedge_dim}, # source_order: (E,) {0: n_node_total_dim + 1}, # source_row_ptr: (N + 1,) + ) + if is_native_spin: + spin = sample_inputs[10] + fparam = sample_inputs[11] + aparam = sample_inputs[12] + charge_spin = sample_inputs[13] + return ( + *base, + # spin: (N, 3) — shares atype's node-axis symbol, same pattern + # as aparam below. + {0: n_node_total_dim} if spin is not None else None, # spin + {0: nframes_dim} if fparam is not None else None, # fparam + {0: n_node_total_dim} if aparam is not None else None, # aparam + {0: nframes_dim} if charge_spin is not None else None, # charge_spin + ) + fparam = sample_inputs[10] + aparam = sample_inputs[11] + charge_spin = sample_inputs[12] + return ( + *base, {0: nframes_dim} if fparam is not None else None, # fparam: (nf, ndf) # aparam: (N, nda) — flat on the node axis, SHARING atype's ``N`` # symbol (the graph fitting consumes aparam per node; an independent @@ -696,6 +833,7 @@ def _build_graph_dynamic_shapes( def _build_graph_dynamic_shapes_with_comm( *sample_inputs: torch.Tensor | None, + is_native_spin: bool = False, ) -> tuple: """Build dynamic-shape specs for the with-comm graph-form export. @@ -716,6 +854,11 @@ def _build_graph_dynamic_shapes_with_comm( source_row_ptr, fparam, aparam, charge_spin, send_list, send_proc, recv_proc, send_num, recv_num, communicator, nlocal, nghost)`` -- 21 entries matching ``forward_lower_graph_exportable_with_comm``. + Native-spin ABI (``is_native_spin=True``): 22 entries, with ``spin`` + inserted at slot 10 and the conditional tail shifted to 11-13, so + the comm block starts at 14. + is_native_spin : bool + Whether ``sample_inputs`` follows the native-spin positional ABI. Returns ------- @@ -723,9 +866,11 @@ def _build_graph_dynamic_shapes_with_comm( Per-input dynamic-shape specs (dicts of ``torch.export.Dim`` or ``None``) in the same order as ``sample_inputs``. """ - fparam = sample_inputs[10] - aparam = sample_inputs[11] - charge_spin = sample_inputs[12] + tail_start = 11 if is_native_spin else 10 + spin = sample_inputs[10] if is_native_spin else None + fparam = sample_inputs[tail_start] + aparam = sample_inputs[tail_start + 1] + charge_spin = sample_inputs[tail_start + 2] nframes_val = 1 n_node_total_dim = torch.export.Dim("n_node_total", min=1) nedge_dim = torch.export.Dim("nedge", min=2) @@ -740,6 +885,13 @@ def _build_graph_dynamic_shapes_with_comm( {0: n_node_total_dim + 1}, # destination_row_ptr: (N + 1,) {0: nedge_dim}, # source_order: (E,) {0: n_node_total_dim + 1}, # source_row_ptr: (N + 1,) + # spin: (N, 3) — EXTENDED node axis, shares atype's symbol; present + # only in the native-spin ABI, where it occupies slot 10. + *( + ({0: n_node_total_dim} if spin is not None else None,) + if is_native_spin + else () + ), {0: nframes_val} if fparam is not None else None, # fparam # aparam: (N, nda) — flat on the SAME extended node axis as atype # (owned prefix + ghost rows). @@ -979,7 +1131,7 @@ def _collect_metadata( # Whether multi-rank LAMMPS needs a second "with-comm" AOTI artifact # (per-layer ghost-feature MPI exchange via deepmd_export::border_op). # The C++ DeepPotPTExpt / DeepSpinPTExpt loaders branch on this flag. - meta["has_comm_artifact"] = _needs_with_comm_artifact(model) + meta["has_comm_artifact"] = _needs_with_comm_artifact(model, lower_kind) # Whether the model's regular .pt2 graph consumes the ``mapping`` # tensor to gather per-layer ghost-atom features from local atoms. @@ -1212,6 +1364,18 @@ def deserialize_to_file( ``metadata.json``. """ lower_kind = _resolve_lower_kind(model_file, data, lower_kind) + if data["model"].get("type") == "native_spin" and lower_kind != "graph": + # Native-spin models implement ONLY the NeighborGraph lower; the + # dense/nlist trace branch does not exist for them. The public freeze + # layer resolves this before calling here (see + # deepmd.pt_expt.entrypoints.main.freeze); this guard pins the + # contract for direct programmatic callers with a clear error instead + # of an opaque trace-time failure. + raise ValueError( + "native-spin models implement only the NeighborGraph lower " + f"(got lower_kind={lower_kind!r}); use lower_kind='graph' with a " + ".pt2 output." + ) # A graph lower deploys the fused inference pipeline. The trace runs at # DP_CUDA_INFER >= 2 so the analytic backward and CSR scatter remain custom # operators, while the per-atom virial is mandatory for the LAMMPS Kokkos @@ -1292,21 +1456,32 @@ def _trace_and_export( target_device = _env.DEVICE - # Detect spin model - is_spin = data["model"].get("type") == "spin_ener" + # Detect spin model. Two flavors share the ``is_spin`` gate below (both + # need the spin-only metadata fields — ``ntypes_spin``/``use_spin`` — + # and the nlist-lower spin ABI probes), but only the NATIVE flavor + # (``native_spin``, ``NativeSpinEnergyModel``) + # rides the graph lower: the virtual-atom flavor (``spin_ener``, + # ``SpinModel``) doubles the atom count and has no graph-lower + # implementation. ``is_native_spin`` distinguishes them at every seam + # below (model rebuild, graph rejection, graph sample-input/dynamic-shape + # ABI, trace call site). + is_native_spin = data["model"].get("type") == "native_spin" + is_spin = is_native_spin or data["model"].get("type") == "spin_ener" # 1. Deserialize model on CPU for make_fx tracing. # make_fx with _allow_non_fake_inputs=True keeps real model parameters; # on CUDA the autograd engine requires CUDA streams for those real # tensors during torch.autograd.grad, but proxy-tensor dispatch doesn't # set streams up → assertion failure. Tracing on CPU avoids this. - if is_spin: + if is_spin and not is_native_spin: from deepmd.pt_expt.model.spin_model import ( SpinModel, ) model = SpinModel.deserialize(data["model"]) else: + # Registry-dispatched (incl. native spin, type "native_spin"): the + # pt_expt BaseModel registry returns this backend's torch class. model = BaseModel.deserialize(data["model"]) model.to("cpu") model.eval() @@ -1337,9 +1512,19 @@ def _trace_and_export( check_graph_trace_torch_version(model) if is_spin: - raise NotImplementedError( - "graph-form .pt2 export is not supported for spin models" - ) + # Only the native spin scheme (NativeSpinEnergyModel: per-local-atom + # spin, no virtual atoms) has a graph-lower export + # (forward_lower_graph_exportable, Task 5) -- and only for the + # regular "graph" kind: "dpa1_canonical" is the compressed-DPA1 + # compact ABI, which native spin never targets. The virtual-atom + # scheme (SpinModel / "spin_ener") has no graph-lower + # implementation at all and stays on the dense (nlist) lower. + if not is_native_spin or lower_kind == "dpa1_canonical": + raise NotImplementedError( + "graph-form .pt2 export supports only the native spin " + "scheme (native_spin); virtual-atom spin models " + "export with the dense lower" + ) # Defense-in-depth: every production caller (freeze entrypoint, # compress, _resolve_lower_kind auto) gates on model_uses_graph_lower # upstream, but a direct programmatic call with lower_kind="graph" @@ -1450,7 +1635,7 @@ def _trace_and_export( ) ensure_comm_registered() - if not _needs_with_comm_artifact(model): + if not _needs_with_comm_artifact(model, lower_kind): raise ValueError( "with_comm_dict=True requested but the model's " "descriptor does not need cross-rank message passing " @@ -1480,6 +1665,7 @@ def _trace_and_export( dtype=torch.float64, edge_dtype=edge_dtype, device=torch.device("cpu"), + want_spin=is_native_spin, ) comm_inputs = _make_comm_sample_inputs( nloc=nlocal_sample, @@ -1495,7 +1681,9 @@ def _trace_and_export( tracing_mode="symbolic", _allow_non_fake_inputs=True, ) - dynamic_shapes = _build_graph_dynamic_shapes_with_comm(*sample_inputs) + dynamic_shapes = _build_graph_dynamic_shapes_with_comm( + *sample_inputs, is_native_spin=is_native_spin + ) else: edge_dtype = ( torch.float32 @@ -1510,18 +1698,39 @@ def _trace_and_export( dtype=torch.float64, edge_dtype=edge_dtype, device=torch.device("cpu"), + want_spin=is_native_spin, ) - traced = model.forward_lower_graph_exportable( - *sample_inputs[:10], - fparam=sample_inputs[10], - aparam=sample_inputs[11], - do_atomic_virial=do_atomic_virial, - charge_spin=sample_inputs[12], - destination_sorted=True, - tracing_mode="symbolic", - _allow_non_fake_inputs=True, + if is_native_spin: + # Native-spin ABI (NativeSpinEnergyModel.forward_lower_graph_exportable): + # slot 10 is ``spin`` (mandatory), slots 11/12 are + # fparam/aparam, slot 13 the conditional ``charge_spin`` tail + # (combined native-spin + charge-spin FiLM models; None + # otherwise). + traced = model.forward_lower_graph_exportable( + *sample_inputs[:10], + spin=sample_inputs[10], + fparam=sample_inputs[11], + aparam=sample_inputs[12], + charge_spin=sample_inputs[13], + do_atomic_virial=do_atomic_virial, + destination_sorted=True, + tracing_mode="symbolic", + _allow_non_fake_inputs=True, + ) + else: + traced = model.forward_lower_graph_exportable( + *sample_inputs[:10], + fparam=sample_inputs[10], + aparam=sample_inputs[11], + do_atomic_virial=do_atomic_virial, + charge_spin=sample_inputs[12], + destination_sorted=True, + tracing_mode="symbolic", + _allow_non_fake_inputs=True, + ) + dynamic_shapes = _build_graph_dynamic_shapes( + *sample_inputs, is_native_spin=is_native_spin ) - dynamic_shapes = _build_graph_dynamic_shapes(*sample_inputs) sample_out = traced(*sample_inputs) output_keys = list(sample_out.keys()) exported = torch.export.export( @@ -1636,7 +1845,7 @@ def _trace_and_export( ) ensure_comm_registered() - if not _needs_with_comm_artifact(model): + if not _needs_with_comm_artifact(model, lower_kind): raise ValueError( "with_comm_dict=True requested but the model's descriptor " "does not need cross-rank message passing " diff --git a/deepmd/pt_expt/utils/vesin_neighbor_list.py b/deepmd/pt_expt/utils/vesin_neighbor_list.py index 3daee763fb..2d28e1b090 100644 --- a/deepmd/pt_expt/utils/vesin_neighbor_list.py +++ b/deepmd/pt_expt/utils/vesin_neighbor_list.py @@ -110,6 +110,8 @@ def build( box_t = box_t.reshape(nframes, 3, 3) if return_mode == "edges": + # LEGACY: serves only the deprecated pt-backend edge_vec .pt2 + # schema; scheduled for removal with that rail. frame_edges = [ _build_single_edges( coord_t[ff], diff --git a/deepmd/tf2/model/dp_model.py b/deepmd/tf2/model/dp_model.py index 4b9258a05c..1326e1c881 100644 --- a/deepmd/tf2/model/dp_model.py +++ b/deepmd/tf2/model/dp_model.py @@ -84,7 +84,10 @@ def call_common( charge_spin: xp.ndarray | None = None, neighbor_list: NeighborList | None = None, ) -> dict[str, xp.ndarray]: - cc, bb, fp, ap, cs, input_prec = self._input_type_cast( + # ``_input_type_cast`` (dpmodel make_model) returns a ``spin`` slot + # for the native-spin graph route; tf2 has no spin/graph lower, so + # it is discarded here. + cc, bb, fp, ap, cs, _, input_prec = self._input_type_cast( to_tensorflow_array(coord), box=to_tensorflow_array(box), fparam=to_tensorflow_array(fparam), @@ -183,7 +186,9 @@ def _call_common_lower_formatted( nlist = to_tensorflow_array(nlist) nframes, _nall = extended_atype.shape[:2] extended_coord = xp.reshape(extended_coord, (nframes, -1, 3)) - cc_ext, _, fp, ap, cs, input_prec = self._input_type_cast( + # ``_input_type_cast`` returns a ``spin`` slot (native-spin graph + # route); tf2 has no spin lower, so it is discarded. + cc_ext, _, fp, ap, cs, _, input_prec = self._input_type_cast( extended_coord, fparam=to_tensorflow_array(fparam), aparam=to_tensorflow_array(aparam), diff --git a/deepmd/tf2/train/trainer.py b/deepmd/tf2/train/trainer.py index 1c1205e1ca..f631c36620 100644 --- a/deepmd/tf2/train/trainer.py +++ b/deepmd/tf2/train/trainer.py @@ -871,7 +871,9 @@ def compiled_prepare_lower_batch( aparam: Any, charge_spin: Any, ) -> tuple[Any, Any, Any, Any, Any, Any, Any, Any, bool]: - cc, bb, fp, ap, cs, _input_prec = model._input_type_cast( + # ``_input_type_cast`` returns a ``spin`` slot (native-spin graph + # route); tf2 has no spin lower, so it is discarded. + cc, bb, fp, ap, cs, _, _input_prec = model._input_type_cast( to_tensorflow_array(coord), box=to_tensorflow_array(box), fparam=to_tensorflow_array(fparam), diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 91dfab4824..2b0d5b9132 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3459,21 +3459,21 @@ def sezm_model_args() -> Argument: str, optional=True, default="None", - doc=doc_only_pt_supported + doc_bridging_method, + doc=doc_bridging_method, ), Argument( "bridging_r_inner", float, optional=True, default=0.5, - doc=doc_only_pt_supported + doc_bridging_r_inner, + doc=doc_bridging_r_inner, ), Argument( "bridging_r_outer", float, optional=True, default=0.8, - doc=doc_only_pt_supported + doc_bridging_r_outer, + doc=doc_bridging_r_outer, ), Argument( "lora", diff --git a/deepmd/utils/spin.py b/deepmd/utils/spin.py index b03cfb07bd..112ad91b3d 100644 --- a/deepmd/utils/spin.py +++ b/deepmd/utils/spin.py @@ -8,6 +8,51 @@ ) +def normalize_spin_use_spin(use_spin: list, type_map: list[str]) -> list[bool]: + """Normalize ``use_spin`` to a per-type boolean list. + + Three equivalent forms are accepted: a per-type boolean list, a list of + magnetic type indices, or a list of magnetic element symbols. The index + and symbol forms are expanded against ``type_map``, so a large type map + only needs its magnetic species named. Pure: the inputs are not + modified. + + Parameters + ---------- + use_spin : list + The ``spin.use_spin`` configuration value, in any of the three + accepted forms. + type_map : list[str] + The model's type map, defining the per-type order and (for the + symbol form) the element names. + + Returns + ------- + list[bool] + The per-type boolean form, of length ``len(type_map)``. + + Raises + ------ + ValueError + If a symbol in ``use_spin`` is absent from ``type_map``. + """ + if use_spin and isinstance(use_spin[0], str): + type_index = {name: idx for idx, name in enumerate(type_map)} + unknown = [name for name in use_spin if name not in type_index] + if unknown: + raise ValueError( + f"spin.use_spin references element(s) {unknown} absent from type_map." + ) + use_spin = [type_index[name] for name in use_spin] + # ``bool`` is a subclass of ``int``; an already-boolean list is passed + # through while an index list is scattered into a per-type mask. + if not use_spin or not isinstance(use_spin[0], bool): + mask = np.full(len(type_map), False, dtype=bool) + mask[use_spin] = True + return mask.tolist() + return [bool(flag) for flag in use_spin] + + class Spin: """Class for spin, mainly processes the spin type-related information. Atom types can be split into three kinds: diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index 94270a8098..cc630b0273 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -432,24 +432,57 @@ pair_coeff * * O H ### Multi-GPU (MPI) inference -The exported `.pt2` runs across multiple GPUs in LAMMPS using MPI domain -decomposition. Multi-GPU support is built into the package by `dp --pt freeze`, -so no extra freeze options are needed and the same `.pt2` file serves both -single- and multi-GPU runs. +:::{important} +**Multi-rank support depends on the model, not on a dense/graph choice.** +DPA4/SeZM reads ghost-neighbour features at every interaction block, so a +multi-rank archive must embed a with-comm AOTInductor artifact +(`model/extra/forward_lower_with_comm.pt2`) for the cross-rank exchange. +Two different export routes produce one: + +- **PT (`dp --pt freeze`, the default on this page).** DPA4/SeZM is detected + and exported through `freeze_sezm_to_pt2`, which uses the compact + **`edge_vec`** ABI (`coord`, `atype`, `edge_index`, `edge_vec`, + `edge_scatter_index`, `edge_mask`), building the neighbour topology on the + C++ side. This is *not* a padded dense neighbour-list export. A plain + energy model reports `supports_edge_parallel() == True`, so the archive + carries the with-comm artifact (`has_comm_artifact=true`) and **supports + multi-rank LAMMPS out of the box** — no extra freeze options. +- **pt_expt (`dp --pt_expt freeze`).** Graph-capable models export through the + **NeighborGraph** ABI (see [Graph-native inference route + (pt_expt)](#graph-native-inference-route-pt_expt) below), which likewise + embeds a with-comm artifact and supports multi-rank LAMMPS. + +The two are distinct ABIs consumed by different C++ paths; neither is a dense +route, and `--lower-kind` selects between the pt_expt lowers only. + +Which models lose multi-rank, on either route: + +- **ZBL zone bridging.** The Source Freeze Propagation gate folds each node's + full *outgoing*-edge set, which no single rank observes for ghost owners, so + `supports_edge_parallel()` is `False` and no with-comm artifact is emitted. + Bridged archives are single-rank; a multi-rank run fails with a clear error + rather than silently dropping the exchange. +- **The deepspin (virtual-atom) spin scheme**, which overrides the export ABI + to `nlist` because it expands virtual atoms inside the graph. + +Native spin (`scheme: "native"`) and charge/spin conditioning are *not* in +that list: native spin reuses the `edge_vec` interface on PT and the +NeighborGraph lower on pt_expt, and both support multi-rank. The remainder of +this subsection describes the multi-GPU launch recipe. +::: -Launch LAMMPS with one MPI rank per GPU and make the target devices visible: +The exported `.pt2` runs across multiple GPUs in LAMMPS using MPI domain +decomposition, with the same `.pt2` file serving both single- and multi-GPU +runs and no extra freeze options needed beyond `--lower-kind graph`. The +launch recipe: ```bash CUDA_VISIBLE_DEVICES=0,1,2,3 mpirun -np 4 lmp -in in.lammps ``` Each MPI rank uses at most one GPU, so `CUDA_VISIBLE_DEVICES` must list every -GPU the run may use. If only one device is visible, all ranks share it: results -stay correct, but the GPU work is serialized and that device's memory grows -with the rank count. - -DPA4/SeZM exchanges neighbor information across the domain boundary, so the -LAMMPS atom map must be enabled: +GPU the run may use. DPA4/SeZM exchanges neighbor information across the +domain boundary, so the LAMMPS atom map must be enabled: ```lammps atom_modify map yes @@ -466,8 +499,117 @@ Two settings improve multi-GPU runs: memory stable. A zero skin rebuilds the neighbor list every step and can substantially increase memory use. -Multi-GPU inference applies to the plain energy model. ZBL zone bridging and -spin models run on a single MPI rank. +## Graph-native inference route (pt_expt) + +In the pt_expt backend, a plain-energy DPA4/SeZM descriptor -- or one +configured with `spin.scheme: native` (see [Native spin +(magnetic)](#native-spin-magnetic) below) -- can be frozen through a +NeighborGraph-native inference path instead of the legacy dense +neighbor-list path. Frame-level charge/spin conditioning +(`add_chg_spin_ebd`) and ZBL zone bridging are graph-eligible too; only the +`deepspin` virtual-atom spin scheme remains dense-only: + +```bash +dp --pt_expt freeze -o model.pt2 --lower-kind graph +``` + +As with DPA-1's and DPA-2's graph paths (see [Difference among different +backends](train-se-atten.md#difference-among-different-backends) and the +"Graph-native inference route (pt_expt)" section of [DPA-2's +documentation](dpa2.md)), the graph route is a sel-free, carry-all builder: +it considers every neighbor within `rcut` rather than the capacity fixed by +`sel`. For the default conservative-energy path, `sel` is already only an +initial search capacity that grows on demand (see the note in [Minimal +input](#minimal-input)), so graph and dense results agree there down to +floating-point noise. Where `sel` does act as a hard cap -- the `dens` +denoising path, or the `deepspin` virtual-atom spin scheme, both of which +size the neighbor list to `sum(sel)` -- the graph route's larger neighbor set +can diverge from the dense route, the same deliberate divergence documented +for DPA-1/DPA-2. + +A DPA4/SeZM descriptor configured with `deepspin`-scheme spin is not +graph-eligible and always runs the dense route regardless of `--lower-kind`; +`--lower-kind graph` on such a model raises an error at freeze time instead +of exporting a silently-dense-only artifact. `native`-scheme spin and ZBL +zone bridging are the opposite case: they have *no* dense route at all -- +the analytical bridging term has no dense injection site -- so they are +always graph-frozen, `--lower-kind` notwithstanding. Charge/spin +conditioning rides the graph lower and constrains neither. Note that a +graph-capable model is always frozen to the graph lower in any case, since +the dense lower is deprecated in the pt_expt backend. See [Native spin (magnetic)](#native-spin-magnetic) below. + +Unlike the dense route (see [Multi-GPU (MPI) +inference](#multi-gpu-mpi-inference) above), a graph-frozen `.pt2` **of a +plain-energy (non-spin) model** embeds a with-comm AOTInductor artifact and +supports multi-rank LAMMPS: each block's cross-rank ghost-feature exchange +runs through the `border_op` MPI path once per interaction block, the same +mechanism used by DPA-2's graph route (see the "Graph-native inference route +(pt_expt)" section of [DPA-2's documentation](dpa2.md)). As on DPA-2, +multi-rank inference on the graph route requires every MPI rank to own or +ghost at least one atom; a rank with zero atoms in both categories aborts the +run collectively rather than silently desynchronizing the per-block +exchange. Pick a domain decomposition that keeps every rank non-empty, or use +the dense route, which has no such restriction (but is single-rank only, as +noted above). + +**Native-spin graph `.pt2` archives are the exception: they carry no +with-comm artifact and are single-rank only** -- see [Native spin +(magnetic)](#native-spin-magnetic) below. + +### Native spin (magnetic) + +`spin.scheme: native` (see [Spin](#spin) above for the general convention and +the `native`-scheme JSON example) is graph-only: a native-spin DPA4/SeZM +model has no dense (nlist) lower at all, so it is always frozen through the +NeighborGraph path regardless of `--lower-kind`, and always runs through it +at inference time. The `deepspin` virtual-atom scheme is unaffected and +keeps using the dense route described in [Spin](#spin) and [Multi-GPU (MPI) +inference](#multi-gpu-mpi-inference). + +- **Native scheme only.** `deepspin`-scheme spin (and the general `spin` + virtual-atom model outside DPA4/SeZM) is dense-only; only `scheme: native` + is graph-eligible. `dp --pt_expt freeze --lower-kind graph` on a + `deepspin`-scheme model raises an error at freeze time, per the dense/graph + eligibility rule above. +- **Graph route only, no dense fallback.** Unlike a plain-energy DPA4/SeZM + descriptor -- which can freeze to either the dense or the graph lower -- + a native-spin descriptor has only the graph lower. `--lower-kind auto` + (the default) resolves to `graph`; `--lower-kind nlist` is not a valid + option for a native-spin model. +- **Single-rank only.** The frozen archive's `has_comm_artifact` metadata is + `false` for native-spin models (no ghost-spin cross-rank exchange is + implemented), so a multi-rank LAMMPS run fails fast with an explicit error + at the first force evaluation, mirroring the dense route's single-rank + restriction described in [Multi-GPU (MPI) + inference](#multi-gpu-mpi-inference). Run native-spin models on a single + MPI rank (a single GPU, or CPU without `mpirun`). +- **Spin is per local atom.** The `spin` input is `(nframes, nloc, 3)` -- + one vector per *local* atom, not per ghost/extended atom (`nall`); there is + no ghost-spin exchange to populate ghost spins across a rank boundary, + consistent with the single-rank restriction above. +- **The magnetic force is a second energy gradient.** As in the general + native-scheme convention (see [Spin](#spin) above), + `force_mag = -\partial E/\partial\mathbf{s}`, computed by pt_expt as a + second `torch.autograd.grad` call alongside the ordinary + `force = -\partial E/\partial\mathbf{r}` call; both are real autograd + outputs, not placeholders. The dpmodel (NumPy) backend is energy-only for + native-spin models -- it has no autograd, so `force`/`force_mag`/`virial` + are `None` placeholders there, exactly as for the plain-energy dpmodel + route. + +The following combinations are **not yet supported** on the native-spin +graph route (follow-up work): + +- **Multi-rank inference.** Ghost-spin cross-rank exchange (analogous to the + plain-energy graph route's `border_op`-based ghost-feature exchange) is not + implemented. +- **Charge-spin FiLM conditioning.** Combining `add_chg_spin_ebd` with + `spin.scheme: native` is rejected at model-construction time; use one or + the other. +- **ZBL zone bridging.** Combining `bridging_method: ZBL` with + `spin.scheme: native` is not supported on the pt_expt backend (`bridging_method` + is rejected there independently of spin -- see [Zone bridging + (ZBL)](#zone-bridging-zbl)). ## Embedding extraction @@ -584,8 +726,24 @@ closed over the one-hop neighbor shell. - DPA4/SeZM is implemented for the PyTorch backend only. - Export uses `.pt2` (AOTInductor); the TorchScript freeze path is not used. - Model compression is not supported. -- Multi-GPU (MPI) LAMMPS inference is supported for the plain energy model; - ZBL zone bridging and spin models run on a single MPI rank. +- Multi-rank (multi-GPU/MPI) LAMMPS inference works for a plain energy model + on both export routes: the PT `edge_vec` archive and the pt_expt + NeighborGraph archive each embed a with-comm artifact. ZBL zone bridging is + single-rank (its Source Freeze Propagation gate folds each node's full + outgoing-edge set, which no single rank observes), and a multi-rank run of + such an archive fails fast rather than dropping the exchange. See + [Multi-GPU (MPI) inference](#multi-gpu-mpi-inference). +- The pt_expt graph-native inference route is unavailable only for + `deepspin`-scheme spin, which stays on the dense route. Charge/spin + conditioning, ZBL zone bridging and `native`-scheme spin are all + graph-eligible. +- `spin.scheme: native` is graph-only (it has no dense route) and supports + multi-rank LAMMPS: ghost node features ride `border_op` per interaction + block and ghost spins arrive through the LAMMPS `sp` forward-comm. + Charge-spin FiLM conditioning combines with it. Combining it with ZBL zone + bridging works single-rank; that combination is single-rank for the same + bridging reason as above. See [Native spin + (magnetic)](#native-spin-magnetic). ## Citation diff --git a/doc/third-party/lammps-command.md b/doc/third-party/lammps-command.md index 94632e4b65..96bd4972ba 100644 --- a/doc/third-party/lammps-command.md +++ b/doc/third-party/lammps-command.md @@ -145,7 +145,7 @@ pair_style deepspin models ... keyword value ... - models = frozen model(s) to compute the interaction. If multiple models are provided, then only the first model serves to provide energy, force and magnetic force prediction for each timestep of molecular dynamics, and the model deviation will be computed among all models every `out_freq` timesteps. -- keyword = _out_file_ or _out_freq_ or _fparam_ or _fparam_from_compute_ or _aparam_from_compute_ or _atomic_ or _relative_ or _aparam_ or _ttm_ +- keyword = _out_file_ or _out_freq_ or _fparam_ or _fparam_from_compute_ or _aparam_from_compute_ or _charge_spin_ or _atomic_ or _relative_ or _aparam_ or _ttm_ :::{note} Please note that the virial and atomic virial are not currently supported in spin models. @@ -162,6 +162,8 @@ Please note that the virial and atomic virial are not currently supported in spi id = compute id used to update the frame parameter. aparam_from_compute value = id id = compute id used to update the atom parameter. + charge_spin value = parameters + parameters = one or more charge/spin condition values required by models trained with a charge/spin embedding. atomic = no value is required. If this keyword is set, the force and magnetic force model deviation of each atom will be output. relative value = level @@ -186,6 +188,8 @@ compute TEMP all temp pair_style deepspin spin.pb aparam_from_compute 1 compute 1 all ke/atom + +pair_style deepspin dpa4_spin.pt2 charge_spin 1.0 2.0 ``` ### Description @@ -194,6 +198,8 @@ Evaluate the interaction of the system with spin by using [DeepSPIN][dpspin] mod This pair style takes the deep spin model defined in a model file that usually has .pb/.pth/.savedmodel extensions. The model can be trained and frozen from multiple backends by package [DeePMD-kit](https://github.com/deepmodeling/deepmd-kit), which can have either double or single float precision interface. +If the keyword `charge_spin` is set, the given per-frame charge/spin value(s) will be fed to models that were trained with a charge/spin embedding. If the keyword is not set, the model's stored `default_chg_spin` (if any) is used. When multiple models are given, the same `charge_spin` is fed to every model. + The model deviation evaluates the consistency of the force and magnetic force predictions from multiple models. By default, only the maximal, minimal and average model deviations are output. If the key `atomic` is set, then the model deviation of force and magnetic force prediction of each atom will be output. The unit follows [LAMMPS units](#units) and the [scale factor](https://docs.lammps.org/pair_hybrid.html) is not applied. diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index 8cc3b715ef..b86ef773bb 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -13,7 +13,7 @@ extern "C" { /** C API version. Bumped whenever the API is changed. * @since API version 22 */ -#define DP_C_API_VERSION 28 +#define DP_C_API_VERSION 29 /** * @brief Neighbor list. @@ -1043,6 +1043,218 @@ extern void DP_DeepSpinComputeNListf2(DP_DeepSpin* dp, float* atomic_energy, float* atomic_virial); +/** + * @brief Evaluate the energy, force, magnetic force and virial by using a DP + *spin model. (double version, with charge_spin) + * @param[in] dp The DP spin model to use. + * @param[in] nframes The number of frames. + * @param[in] natoms The number of atoms. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] spin The spins of atoms. The array should be of size nframes x + *natoms x 3. + * @param[in] atype The atom types. The array should contain natoms ints. + * @param[in] cell The cell of the region. The array should be of size nframes + *x 9. Pass NULL if pbc is not used. + * @param[in] fparam The frame parameters. The array can be of size nframes x + *dim_fparam. + * @param[in] aparam The atom parameters. The array can be of size nframes x + *natoms x dim_aparam. + * @param[in] charge_spin The per-frame charge/spin input. The array can be of + *size nframes x dim_chg_spin. Pass NULL to use the model's stored + *default_chg_spin. + * @param[out] energy Output energy. + * @param[out] force Output force. The array should be of size natoms x 3. + * @param[out] force_mag Output magnetic force. The array should be of size + *natoms x 3. + * @param[out] virial Output virial. The array should be of size 9. + * @param[out] atomic_energy Output atomic energy. The array should be of size + *natoms. + * @param[out] atomic_virial Output atomic virial. The array should be of size + *natoms x 9. + * @warning The output arrays should be allocated before calling this function. + *Pass NULL if not required. + * @since API version 29 + **/ +extern void DP_DeepSpinCompute3(DP_DeepSpin* dp, + const int nframes, + const int natoms, + const double* coord, + const double* spin, + const int* atype, + const double* cell, + const double* fparam, + const double* aparam, + const double* charge_spin, + double* energy, + double* force, + double* force_mag, + double* virial, + double* atomic_energy, + double* atomic_virial); + +/** + * @brief Evaluate the energy, force, magnetic force and virial by using a DP + *spin model. (float version, with charge_spin) + * @param[in] dp The DP spin model to use. + * @param[in] nframes The number of frames. + * @param[in] natoms The number of atoms. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] spin The spins of atoms. The array should be of size nframes x + *natoms x 3. + * @param[in] atype The atom types. The array should contain natoms ints. + * @param[in] cell The cell of the region. The array should be of size nframes + *x 9. Pass NULL if pbc is not used. + * @param[in] fparam The frame parameters. The array can be of size nframes x + *dim_fparam. + * @param[in] aparam The atom parameters. The array can be of size nframes x + *natoms x dim_aparam. + * @param[in] charge_spin The per-frame charge/spin input. The array can be of + *size nframes x dim_chg_spin. Pass NULL to use the model's stored + *default_chg_spin. + * @param[out] energy Output energy. + * @param[out] force Output force. The array should be of size natoms x 3. + * @param[out] force_mag Output magnetic force. The array should be of size + *natoms x 3. + * @param[out] virial Output virial. The array should be of size 9. + * @param[out] atomic_energy Output atomic energy. The array should be of size + *natoms. + * @param[out] atomic_virial Output atomic virial. The array should be of size + *natoms x 9. + * @warning The output arrays should be allocated before calling this function. + *Pass NULL if not required. + * @since API version 29 + **/ +extern void DP_DeepSpinComputef3(DP_DeepSpin* dp, + const int nframes, + const int natoms, + const float* coord, + const float* spin, + const int* atype, + const float* cell, + const float* fparam, + const float* aparam, + const float* charge_spin, + double* energy, + float* force, + float* force_mag, + float* virial, + float* atomic_energy, + float* atomic_virial); + +/** + * @brief Evaluate the energy, force, magnetic force and virial by using a DP + *spin model with the neighbor list. (double version, with charge_spin) + * @param[in] dp The DP spin model to use. + * @param[in] nframes The number of frames. + * @param[in] natoms The number of atoms. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] spin The spins of atoms. The array should be of size nframes x + *natoms x 3. + * @param[in] atype The atom types. The array should contain natoms ints. + * @param[in] cell The cell of the region. The array should be of size nframes + *x 9. Pass NULL if pbc is not used. + * @param[in] nghost The number of ghost atoms. + * @param[in] nlist The neighbor list. + * @param[in] ago Update the internal neighbour list if ago is 0. + * @param[in] fparam The frame parameters. The array can be of size nframes x + *dim_fparam. + * @param[in] aparam The atom parameters. The array can be of size nframes x + *natoms x dim_aparam. + * @param[in] charge_spin The per-frame charge/spin input. The array can be of + *size nframes x dim_chg_spin. Pass NULL to use the model's stored + *default_chg_spin. + * @param[out] energy Output energy. + * @param[out] force Output force. The array should be of size natoms x 3. + * @param[out] force_mag Output magnetic force. The array should be of size + *natoms x 3. + * @param[out] virial Output virial. The array should be of size 9. + * @param[out] atomic_energy Output atomic energy. The array should be of size + *natoms. + * @param[out] atomic_virial Output atomic virial. The array should be of size + *natoms x 9. + * @warning The output arrays should be allocated before calling this function. + *Pass NULL if not required. + * @since API version 29 + **/ +extern void DP_DeepSpinComputeNList3(DP_DeepSpin* dp, + const int nframes, + const int natoms, + const double* coord, + const double* spin, + const int* atype, + const double* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const double* fparam, + const double* aparam, + const double* charge_spin, + double* energy, + double* force, + double* force_mag, + double* virial, + double* atomic_energy, + double* atomic_virial); + +/** + * @brief Evaluate the energy, force, magnetic force and virial by using a DP + *spin model with the neighbor list. (float version, with charge_spin) + * @param[in] dp The DP spin model to use. + * @param[in] nframes The number of frames. + * @param[in] natoms The number of atoms. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] spin The spins of atoms. The array should be of size nframes x + *natoms x 3. + * @param[in] atype The atom types. The array should contain natoms ints. + * @param[in] cell The cell of the region. The array should be of size nframes + *x 9. Pass NULL if pbc is not used. + * @param[in] nghost The number of ghost atoms. + * @param[in] nlist The neighbor list. + * @param[in] ago Update the internal neighbour list if ago is 0. + * @param[in] fparam The frame parameters. The array can be of size nframes x + *dim_fparam. + * @param[in] aparam The atom parameters. The array can be of size nframes x + *natoms x dim_aparam. + * @param[in] charge_spin The per-frame charge/spin input. The array can be of + *size nframes x dim_chg_spin. Pass NULL to use the model's stored + *default_chg_spin. + * @param[out] energy Output energy. + * @param[out] force Output force. The array should be of size natoms x 3. + * @param[out] force_mag Output magnetic force. The array should be of size + *natoms x 3. + * @param[out] virial Output virial. The array should be of size 9. + * @param[out] atomic_energy Output atomic energy. The array should be of size + *natoms. + * @param[out] atomic_virial Output atomic virial. The array should be of size + *natoms x 9. + * @warning The output arrays should be allocated before calling this function. + *Pass NULL if not required. + * @since API version 29 + **/ +extern void DP_DeepSpinComputeNListf3(DP_DeepSpin* dp, + const int nframes, + const int natoms, + const float* coord, + const float* spin, + const int* atype, + const float* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const float* fparam, + const float* aparam, + const float* charge_spin, + double* energy, + float* force, + float* force_mag, + float* virial, + float* atomic_energy, + float* atomic_virial); + /** * @brief Evaluate the energy, force and virial by using a DP with the mixed *type. (double version) @@ -1926,6 +2138,234 @@ void DP_DeepSpinModelDeviComputeNListf2(DP_DeepSpinModelDevi* dp, float* atomic_energy, float* atomic_virial); +/** + * @brief Evaluate the energy, force, magnetic force and virial by using a DP + *spin model deviation. (double version, with charge_spin) + * @version 3 + * @param[in] dp The DP model deviation to use. + * @param[in] nframes The number of frames. Only support 1 for now. + * @param[in] natoms The number of atoms. + * @param[in] coord The coordinates of atoms. The array should be of size natoms + *x 3. + * @param[in] spin The spins of atoms, [0, 0, 0] if no spin. The array should be + *of size nframes x natoms x 3. + * @param[in] atype The atom types. The array should contain natoms ints. + * @param[in] cell The cell of the region. The array should be of size 9. Pass + *NULL if pbc is not used. + * @param[in] fparam The frame parameters. The array can be of size nframes x + *dim_fparam. + * @param[in] aparam The atom parameters. The array can be of size nframes x + *natoms x dim_aparam. + * @param[in] charge_spin The per-frame charge/spin input. The array can be of + *size nframes x dim_chg_spin. Pass NULL to use the model's stored + *default_chg_spin. + * @param[out] energy Output energies of all models. The array should be of size + *nmodels. + * @param[out] force Output forces of all models. The array should be of size + *nmodels x natoms x 3. + * @param[out] force_mag Output magnetic forces of all models. The array should + *be of size nmodels x natoms x 3. + * @param[out] virial Output virials of all models. The array should be of size + *nmodels x 9. + * @param[out] atomic_energy Output atomic energies of all models. The array + *should be of size nmodels x natoms. + * @param[out] atomic_virial Output atomic virials of all models. The array + *should be of size nmodels x natoms x 9. + * @warning The output arrays should be allocated before calling this function. + *Pass NULL if not required. + * @since API version 29 + **/ +void DP_DeepSpinModelDeviCompute3(DP_DeepSpinModelDevi* dp, + const int nframes, + const int natoms, + const double* coord, + const double* spin, + const int* atype, + const double* cell, + const double* fparam, + const double* aparam, + const double* charge_spin, + double* energy, + double* force, + double* force_mag, + double* virial, + double* atomic_energy, + double* atomic_virial); + +/** + * @brief Evaluate the energy, force, magnetic force and virial by using a DP + *spin model deviation. (float version, with charge_spin) + * @version 3 + * @param[in] dp The DP model deviation to use. + * @param[in] nframes The number of frames. Only support 1 for now. + * @param[in] natoms The number of atoms. + * @param[in] coord The coordinates of atoms. The array should be of size natoms + *x 3. + * @param[in] spin The spins of atoms, [0, 0, 0] if no spin. The array should be + *of size nframes x natoms x 3. + * @param[in] atype The atom types. The array should contain natoms ints. + * @param[in] cell The cell of the region. The array should be of size 9. Pass + *NULL if pbc is not used. + * @param[in] fparam The frame parameters. The array can be of size nframes x + *dim_fparam. + * @param[in] aparam The atom parameters. The array can be of size nframes x + *natoms x dim_aparam. + * @param[in] charge_spin The per-frame charge/spin input. The array can be of + *size nframes x dim_chg_spin. Pass NULL to use the model's stored + *default_chg_spin. + * @param[out] energy Output energies of all models. The array should be of size + *nmodels. + * @param[out] force Output forces of all models. The array should be of size + *nmodels x natoms x 3. + * @param[out] force_mag Output magnetic forces of all models. The array should + *be of size nmodels x natoms x 3. + * @param[out] virial Output virials of all models. The array should be of size + *nmodels x 9. + * @param[out] atomic_energy Output atomic energies of all models. The array + *should be of size nmodels x natoms. + * @param[out] atomic_virial Output atomic virials of all models. The array + *should be of size nmodels x natoms x 9. + * @warning The output arrays should be allocated before calling this function. + *Pass NULL if not required. + * @since API version 29 + **/ +void DP_DeepSpinModelDeviComputef3(DP_DeepSpinModelDevi* dp, + const int nframes, + const int natoms, + const float* coord, + const float* spin, + const int* atype, + const float* cell, + const float* fparam, + const float* aparam, + const float* charge_spin, + double* energy, + float* force, + float* force_mag, + float* virial, + float* atomic_energy, + float* atomic_virial); + +/** + * @brief Evaluate the energy, force, magnetic force and virial by using a DP + *spin model deviation with neighbor list. (double version, with charge_spin) + * @version 3 + * @param[in] dp The DP model deviation to use. + * @param[in] nframes The number of frames. Only support 1 for now. + * @param[in] natoms The number of atoms. + * @param[in] coord The coordinates of atoms. The array should be of size natoms + *x 3. + * @param[in] spin The spins of atoms, [0, 0, 0] if no spin. The array should be + *of size nframes x natoms x 3. + * @param[in] atype The atom types. The array should contain natoms ints. + * @param[in] cell The cell of the region. The array should be of size 9. Pass + *NULL if pbc is not used. + * @param[in] nghost The number of ghost atoms. + * @param[in] nlist The neighbor list. + * @param[in] ago Update the internal neighbour list if ago is 0. + * @param[in] fparam The frame parameters. The array can be of size nframes x + *dim_fparam. + * @param[in] aparam The atom parameters. The array can be of size nframes x + *natoms x dim_aparam. + * @param[in] charge_spin The per-frame charge/spin input. The array can be of + *size nframes x dim_chg_spin. Pass NULL to use the model's stored + *default_chg_spin. + * @param[out] energy Output energies of all models. The array should be of size + *nmodels. + * @param[out] force Output forces of all models. The array should be of size + *nmodels x natoms x 3. + * @param[out] force_mag Output magnetic forces of all models. The array should + *be of size nmodels x natoms x 3. + * @param[out] virial Output virials of all models. The array should be of size + *nmodels x 9. + * @param[out] atomic_energy Output atomic energies of all models. The array + *should be of size nmodels x natoms. + * @param[out] atomic_virial Output atomic virials of all models. The array + *should be of size nmodels x natoms x 9. + * @warning The output arrays should be allocated before calling this function. + *Pass NULL if not required. + * @since API version 29 + **/ +void DP_DeepSpinModelDeviComputeNList3(DP_DeepSpinModelDevi* dp, + const int nframes, + const int natoms, + const double* coord, + const double* spin, + const int* atype, + const double* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const double* fparam, + const double* aparam, + const double* charge_spin, + double* energy, + double* force, + double* force_mag, + double* virial, + double* atomic_energy, + double* atomic_virial); + +/** + * @brief Evaluate the energy, force, magnetic force and virial by using a DP + *spin model deviation with neighbor list. (float version, with charge_spin) + * @version 3 + * @param[in] dp The DP model deviation to use. + * @param[in] nframes The number of frames. Only support 1 for now. + * @param[in] natoms The number of atoms. + * @param[in] coord The coordinates of atoms. The array should be of size natoms + *x 3. + * @param[in] spin The spins of atoms, [0, 0, 0] if no spin. The array should be + *of size nframes x natoms x 3. + * @param[in] atype The atom types. The array should contain natoms ints. + * @param[in] cell The cell of the region. The array should be of size 9. Pass + *NULL if pbc is not used. + * @param[in] nghost The number of ghost atoms. + * @param[in] nlist The neighbor list. + * @param[in] ago Update the internal neighbour list if ago is 0. + * @param[in] fparam The frame parameters. The array can be of size nframes x + *dim_fparam. + * @param[in] aparam The atom parameters. The array can be of size nframes x + *natoms x dim_aparam. + * @param[in] charge_spin The per-frame charge/spin input. The array can be of + *size nframes x dim_chg_spin. Pass NULL to use the model's stored + *default_chg_spin. + * @param[out] energy Output energies of all models. The array should be of size + *nmodels. + * @param[out] force Output forces of all models. The array should be of size + *nmodels x natoms x 3. + * @param[out] force_mag Output magnetic forces of all models. The array should + *be of size nmodels x natoms x 3. + * @param[out] virial Output virials of all models. The array should be of size + *nmodels x 9. + * @param[out] atomic_energy Output atomic energies of all models. The array + *should be of size nmodels x natoms. + * @param[out] atomic_virial Output atomic virials of all models. The array + *should be of size nmodels x natoms x 9. + * @warning The output arrays should be allocated before calling this function. + *Pass NULL if not required. + * @since API version 29 + **/ +void DP_DeepSpinModelDeviComputeNListf3(DP_DeepSpinModelDevi* dp, + const int nframes, + const int natoms, + const float* coord, + const float* spin, + const int* atype, + const float* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const float* fparam, + const float* aparam, + const float* charge_spin, + double* energy, + float* force, + float* force_mag, + float* virial, + float* atomic_energy, + float* atomic_virial); + // Deep Base Model methods /** * @brief Get the cutoff of a DP. @@ -2263,6 +2703,15 @@ int DP_DeepSpinGetDimFParam(DP_DeepSpin* dp); */ int DP_DeepSpinGetDimAParam(DP_DeepSpin* dp); +/** + * @brief Get the dimension of the charge/spin input of a DP Spin Model. + * @param[in] dp The DP Spin Model to use. + * @return The dimension of the charge/spin input (0 if the model has no + * charge/spin embedding). + * @since API version 29 + */ +int DP_DeepSpinGetDimChgSpin(DP_DeepSpin* dp); + /** * @brief Check whether the atomic dimension of atomic parameters is nall * instead of nloc. @@ -2315,6 +2764,15 @@ int DP_DeepSpinModelDeviGetDimFParam(DP_DeepSpinModelDevi* dp); */ int DP_DeepSpinModelDeviGetDimAParam(DP_DeepSpinModelDevi* dp); +/** + * @brief Get the dimension of the charge/spin input of a DP Spin Model + * Deviation. + * @param[in] dp The DP Spin Model Deviation to use. + * @return The dimension of the charge/spin input (0 if none). + * @since API version 29 + */ +int DP_DeepSpinModelDeviGetDimChgSpin(DP_DeepSpinModelDevi* dp); + /** * @brief Check whether the atomic dimension of atomic parameters is nall * instead of nloc. diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index d8d67dd78b..08cabbcf0c 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -123,6 +123,7 @@ inline void _DP_DeepSpinCompute(DP_DeepSpin* dp, const FPTYPE* cell, const FPTYPE* fparam, const FPTYPE* aparam, + const FPTYPE* charge_spin, double* energy, FPTYPE* force, FPTYPE* force_mag, @@ -140,15 +141,24 @@ inline void _DP_DeepSpinCompute(DP_DeepSpin* dp, const double* cell, const double* fparam, const double* aparam, + const double* charge_spin, double* energy, double* force, double* force_mag, double* virial, double* atomic_energy, double* atomic_virial) { - DP_DeepSpinCompute2(dp, nframes, natom, coord, spin, atype, cell, fparam, - aparam, energy, force, force_mag, virial, atomic_energy, - atomic_virial); + // charge_spin == nullptr keeps the version-2 entry point so models without a + // charge/spin embedding still work against an older libdeepmd_c. + if (charge_spin) { + DP_DeepSpinCompute3(dp, nframes, natom, coord, spin, atype, cell, fparam, + aparam, charge_spin, energy, force, force_mag, virial, + atomic_energy, atomic_virial); + } else { + DP_DeepSpinCompute2(dp, nframes, natom, coord, spin, atype, cell, fparam, + aparam, energy, force, force_mag, virial, atomic_energy, + atomic_virial); + } } template <> @@ -161,15 +171,22 @@ inline void _DP_DeepSpinCompute(DP_DeepSpin* dp, const float* cell, const float* fparam, const float* aparam, + const float* charge_spin, double* energy, float* force, float* force_mag, float* virial, float* atomic_energy, float* atomic_virial) { - DP_DeepSpinComputef2(dp, nframes, natom, coord, spin, atype, cell, fparam, - aparam, energy, force, force_mag, virial, atomic_energy, - atomic_virial); + if (charge_spin) { + DP_DeepSpinComputef3(dp, nframes, natom, coord, spin, atype, cell, fparam, + aparam, charge_spin, energy, force, force_mag, virial, + atomic_energy, atomic_virial); + } else { + DP_DeepSpinComputef2(dp, nframes, natom, coord, spin, atype, cell, fparam, + aparam, energy, force, force_mag, virial, + atomic_energy, atomic_virial); + } } template @@ -263,6 +280,7 @@ inline void _DP_DeepSpinComputeNList(DP_DeepSpin* dp, const int ago, const FPTYPE* fparam, const FPTYPE* aparam, + const FPTYPE* charge_spin, double* energy, FPTYPE* force, FPTYPE* force_mag, @@ -283,15 +301,23 @@ inline void _DP_DeepSpinComputeNList(DP_DeepSpin* dp, const int ago, const double* fparam, const double* aparam, + const double* charge_spin, double* energy, double* force, double* force_mag, double* virial, double* atomic_energy, double* atomic_virial) { - DP_DeepSpinComputeNList2(dp, nframes, natom, coord, spin, atype, cell, nghost, - nlist, ago, fparam, aparam, energy, force, force_mag, - virial, atomic_energy, atomic_virial); + if (charge_spin) { + DP_DeepSpinComputeNList3(dp, nframes, natom, coord, spin, atype, cell, + nghost, nlist, ago, fparam, aparam, charge_spin, + energy, force, force_mag, virial, atomic_energy, + atomic_virial); + } else { + DP_DeepSpinComputeNList2(dp, nframes, natom, coord, spin, atype, cell, + nghost, nlist, ago, fparam, aparam, energy, force, + force_mag, virial, atomic_energy, atomic_virial); + } } template <> @@ -307,15 +333,23 @@ inline void _DP_DeepSpinComputeNList(DP_DeepSpin* dp, const int ago, const float* fparam, const float* aparam, + const float* charge_spin, double* energy, float* force, float* force_mag, float* virial, float* atomic_energy, float* atomic_virial) { - DP_DeepSpinComputeNListf2(dp, nframes, natom, coord, spin, atype, cell, - nghost, nlist, ago, fparam, aparam, energy, force, - force_mag, virial, atomic_energy, atomic_virial); + if (charge_spin) { + DP_DeepSpinComputeNListf3(dp, nframes, natom, coord, spin, atype, cell, + nghost, nlist, ago, fparam, aparam, charge_spin, + energy, force, force_mag, virial, atomic_energy, + atomic_virial); + } else { + DP_DeepSpinComputeNListf2(dp, nframes, natom, coord, spin, atype, cell, + nghost, nlist, ago, fparam, aparam, energy, force, + force_mag, virial, atomic_energy, atomic_virial); + } } template @@ -445,6 +479,7 @@ inline void _DP_DeepSpinModelDeviCompute(DP_DeepSpinModelDevi* dp, const FPTYPE* cell, const FPTYPE* fparam, const FPTYPE* aparam, + const FPTYPE* charge_spin, double* energy, FPTYPE* force, FPTYPE* force_mag, @@ -461,15 +496,24 @@ inline void _DP_DeepSpinModelDeviCompute(DP_DeepSpinModelDevi* dp, const double* cell, const double* fparam, const double* aparam, + const double* charge_spin, double* energy, double* force, double* force_mag, double* virial, double* atomic_energy, double* atomic_virial) { - DP_DeepSpinModelDeviCompute2(dp, 1, natom, coord, spin, atype, cell, fparam, - aparam, energy, force, force_mag, virial, - atomic_energy, atomic_virial); + // charge_spin == nullptr keeps the version-2 entry point so models without a + // charge/spin embedding still work against an older libdeepmd_c. + if (charge_spin) { + DP_DeepSpinModelDeviCompute3(dp, 1, natom, coord, spin, atype, cell, fparam, + aparam, charge_spin, energy, force, force_mag, + virial, atomic_energy, atomic_virial); + } else { + DP_DeepSpinModelDeviCompute2(dp, 1, natom, coord, spin, atype, cell, fparam, + aparam, energy, force, force_mag, virial, + atomic_energy, atomic_virial); + } } template <> @@ -481,15 +525,22 @@ inline void _DP_DeepSpinModelDeviCompute(DP_DeepSpinModelDevi* dp, const float* cell, const float* fparam, const float* aparam, + const float* charge_spin, double* energy, float* force, float* force_mag, float* virial, float* atomic_energy, float* atomic_virial) { - DP_DeepSpinModelDeviComputef2(dp, 1, natom, coord, spin, atype, cell, fparam, - aparam, energy, force, force_mag, virial, - atomic_energy, atomic_virial); + if (charge_spin) { + DP_DeepSpinModelDeviComputef3( + dp, 1, natom, coord, spin, atype, cell, fparam, aparam, charge_spin, + energy, force, force_mag, virial, atomic_energy, atomic_virial); + } else { + DP_DeepSpinModelDeviComputef2(dp, 1, natom, coord, spin, atype, cell, + fparam, aparam, energy, force, force_mag, + virial, atomic_energy, atomic_virial); + } } template @@ -578,6 +629,7 @@ inline void _DP_DeepSpinModelDeviComputeNList(DP_DeepSpinModelDevi* dp, const int ago, const FPTYPE* fparam, const FPTYPE* aparam, + const FPTYPE* charge_spin, double* energy, FPTYPE* force, FPTYPE* force_mag, @@ -596,15 +648,23 @@ inline void _DP_DeepSpinModelDeviComputeNList(DP_DeepSpinModelDevi* dp, const int ago, const double* fparam, const double* aparam, + const double* charge_spin, double* energy, double* force, double* force_mag, double* virial, double* atomic_energy, double* atomic_virial) { - DP_DeepSpinModelDeviComputeNList2( - dp, 1, natom, coord, spin, atype, cell, nghost, nlist, ago, fparam, - aparam, energy, force, force_mag, virial, atomic_energy, atomic_virial); + if (charge_spin) { + DP_DeepSpinModelDeviComputeNList3(dp, 1, natom, coord, spin, atype, cell, + nghost, nlist, ago, fparam, aparam, + charge_spin, energy, force, force_mag, + virial, atomic_energy, atomic_virial); + } else { + DP_DeepSpinModelDeviComputeNList2( + dp, 1, natom, coord, spin, atype, cell, nghost, nlist, ago, fparam, + aparam, energy, force, force_mag, virial, atomic_energy, atomic_virial); + } } template <> inline void _DP_DeepSpinModelDeviComputeNList(DP_DeepSpinModelDevi* dp, @@ -618,15 +678,23 @@ inline void _DP_DeepSpinModelDeviComputeNList(DP_DeepSpinModelDevi* dp, const int ago, const float* fparam, const float* aparam, + const float* charge_spin, double* energy, float* force, float* force_mag, float* virial, float* atomic_energy, float* atomic_virial) { - DP_DeepSpinModelDeviComputeNListf2( - dp, 1, natom, coord, spin, atype, cell, nghost, nlist, ago, fparam, - aparam, energy, force, force_mag, virial, atomic_energy, atomic_virial); + if (charge_spin) { + DP_DeepSpinModelDeviComputeNListf3(dp, 1, natom, coord, spin, atype, cell, + nghost, nlist, ago, fparam, aparam, + charge_spin, energy, force, force_mag, + virial, atomic_energy, atomic_virial); + } else { + DP_DeepSpinModelDeviComputeNListf2( + dp, 1, natom, coord, spin, atype, cell, nghost, nlist, ago, fparam, + aparam, energy, force, force_mag, virial, atomic_energy, atomic_virial); + } } template @@ -1763,7 +1831,7 @@ class DeepSpin : public DeepBaseModel { /** * @brief DP constructor without initialization. **/ - DeepSpin() : dp(nullptr) {}; + DeepSpin() : dp(nullptr), dchgspin(0) {}; ~DeepSpin() { DP_DeleteDeepSpin(dp); }; /** * @brief DP constructor with initialization. @@ -1774,7 +1842,7 @@ class DeepSpin : public DeepBaseModel { DeepSpin(const std::string& model, const int& gpu_rank = 0, const std::string& file_content = "") - : dp(nullptr) { + : dp(nullptr), dchgspin(0) { try { init(model, gpu_rank, file_content); } catch (...) { @@ -1805,11 +1873,22 @@ class DeepSpin : public DeepBaseModel { DP_CHECK_OK(DP_DeepSpinCheckOK, dp); dfparam = DP_DeepSpinGetDimFParam(dp); daparam = DP_DeepSpinGetDimAParam(dp); + dchgspin = DP_DeepSpinGetDimChgSpin(dp); aparam_nall = DP_DeepSpinIsAParamNAll(dp); has_default_fparam_ = DP_DeepSpinHasDefaultFParam(dp); dpbase = (DP_DeepBaseModel*)dp; }; + /** + * @brief Get the dimension of the charge/spin embedding input. + * @return The dimension of the charge/spin input (0 if the model has no + *charge/spin embedding). + **/ + int dim_chg_spin() const { + assert(dp); + return dchgspin; + } + /** * @brief Evaluate the energy, force, magnetic force and virial by using this *DP spin model. @@ -1832,6 +1911,10 @@ class DeepSpin : public DeepBaseModel { * nframes x natoms x dim_aparam. * natoms x dim_aparam. Then all frames are assumed to be provided with the *same aparam. + * @param[in] charge_spin The charge/spin input. The array can be of size: + * nframes x dim_chg_spin. + * dim_chg_spin. Then all frames are assumed to be provided with the same + *charge_spin. Leave it empty to use the model's stored default_chg_spin. * @warning Natoms should not be zero when computing multiple frames. **/ template @@ -1845,7 +1928,8 @@ class DeepSpin : public DeepBaseModel { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()) { + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -1869,10 +1953,15 @@ class DeepSpin : public DeepBaseModel { tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + // charge_spin routes to the version-3 C API; nullptr keeps version-2 so + // non-charge_spin models still work against an older libdeepmd_c. + std::vector charge_spin_tiled_; + const VALUETYPE* charge_spin__ = validate_charge_spin( + charge_spin, dchgspin, nframes, charge_spin_tiled_); - _DP_DeepSpinCompute(dp, nframes, natoms, coord_, spin_, atype_, - box_, fparam__, aparam__, ener_, force_, - force_mag_, virial_, nullptr, nullptr); + _DP_DeepSpinCompute( + dp, nframes, natoms, coord_, spin_, atype_, box_, fparam__, aparam__, + charge_spin__, ener_, force_, force_mag_, virial_, nullptr, nullptr); DP_CHECK_OK(DP_DeepSpinCheckOK, dp); }; @@ -1900,6 +1989,10 @@ class DeepSpin : public DeepBaseModel { * nframes x natoms x dim_aparam. * natoms x dim_aparam. Then all frames are assumed to be provided with the *same aparam. + * @param[in] charge_spin The charge/spin input. The array can be of size: + * nframes x dim_chg_spin. + * dim_chg_spin. Then all frames are assumed to be provided with the same + *charge_spin. Leave it empty to use the model's stored default_chg_spin. * @warning Natoms should not be zero when computing multiple frames. **/ template @@ -1915,7 +2008,8 @@ class DeepSpin : public DeepBaseModel { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()) { + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -1944,10 +2038,16 @@ class DeepSpin : public DeepBaseModel { tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + // charge_spin routes to the version-3 C API; nullptr keeps version-2 so + // non-charge_spin models still work against an older libdeepmd_c. + std::vector charge_spin_tiled_; + const VALUETYPE* charge_spin__ = validate_charge_spin( + charge_spin, dchgspin, nframes, charge_spin_tiled_); - _DP_DeepSpinCompute( - dp, nframes, natoms, coord_, spin_, atype_, box_, fparam__, aparam__, - ener_, force_, force_mag_, virial_, atomic_ener_, atomic_virial_); + _DP_DeepSpinCompute(dp, nframes, natoms, coord_, spin_, atype_, + box_, fparam__, aparam__, charge_spin__, + ener_, force_, force_mag_, virial_, + atomic_ener_, atomic_virial_); DP_CHECK_OK(DP_DeepSpinCheckOK, dp); }; @@ -1976,6 +2076,10 @@ class DeepSpin : public DeepBaseModel { * nframes x natoms x dim_aparam. * natoms x dim_aparam. Then all frames are assumed to be provided with the *same aparam. + * @param[in] charge_spin The charge/spin input. The array can be of size: + * nframes x dim_chg_spin. + * dim_chg_spin. Then all frames are assumed to be provided with the same + *charge_spin. Leave it empty to use the model's stored default_chg_spin. * @warning Natoms should not be zero when computing multiple frames. **/ template @@ -1992,7 +2096,8 @@ class DeepSpin : public DeepBaseModel { const InputNlist& lmp_list, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()) { + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -2019,10 +2124,15 @@ class DeepSpin : public DeepBaseModel { aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; - _DP_DeepSpinComputeNList(dp, nframes, natoms, coord_, spin_, - atype_, box_, nghost, lmp_list.nl, ago, - fparam__, aparam__, ener_, force_, - force_mag_, virial_, nullptr, nullptr); + // charge_spin routes to the version-3 C API; nullptr keeps version-2 so + // non-charge_spin models still work against an older libdeepmd_c. + std::vector charge_spin_tiled_; + const VALUETYPE* charge_spin__ = validate_charge_spin( + charge_spin, dchgspin, nframes, charge_spin_tiled_); + _DP_DeepSpinComputeNList( + dp, nframes, natoms, coord_, spin_, atype_, box_, nghost, lmp_list.nl, + ago, fparam__, aparam__, charge_spin__, ener_, force_, force_mag_, + virial_, nullptr, nullptr); DP_CHECK_OK(DP_DeepSpinCheckOK, dp); }; @@ -2053,6 +2163,10 @@ class DeepSpin : public DeepBaseModel { * nframes x natoms x dim_aparam. * natoms x dim_aparam. Then all frames are assumed to be provided with the *same aparam. + * @param[in] charge_spin The charge/spin input. The array can be of size: + * nframes x dim_chg_spin. + * dim_chg_spin. Then all frames are assumed to be provided with the same + *charge_spin. Leave it empty to use the model's stored default_chg_spin. * @warning Natoms should not be zero when computing multiple frames. **/ template @@ -2071,7 +2185,8 @@ class DeepSpin : public DeepBaseModel { const InputNlist& lmp_list, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()) { + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -2102,15 +2217,21 @@ class DeepSpin : public DeepBaseModel { aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + // charge_spin routes to the version-3 C API; nullptr keeps version-2 so + // non-charge_spin models still work against an older libdeepmd_c. + std::vector charge_spin_tiled_; + const VALUETYPE* charge_spin__ = validate_charge_spin( + charge_spin, dchgspin, nframes, charge_spin_tiled_); _DP_DeepSpinComputeNList( dp, nframes, natoms, coord_, spin_, atype_, box_, nghost, lmp_list.nl, - ago, fparam__, aparam__, ener_, force_, force_mag_, virial_, - atomic_ener_, atomic_virial_); + ago, fparam__, aparam__, charge_spin__, ener_, force_, force_mag_, + virial_, atomic_ener_, atomic_virial_); DP_CHECK_OK(DP_DeepSpinCheckOK, dp); }; private: DP_DeepSpin* dp; + int dchgspin; }; /** @@ -2827,13 +2948,14 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { /** * @brief DP model deviation constructor without initialization. **/ - DeepSpinModelDevi() : dp(nullptr) {}; + DeepSpinModelDevi() : dp(nullptr), dchgspin(0) {}; ~DeepSpinModelDevi() { DP_DeleteDeepSpinModelDevi(dp); }; /** * @brief DP model deviation constructor with initialization. * @param[in] models The names of the frozen model file. **/ - DeepSpinModelDevi(const std::vector& models) : dp(nullptr) { + DeepSpinModelDevi(const std::vector& models) + : dp(nullptr), dchgspin(0) { try { init(models); } catch (...) { @@ -2882,11 +3004,22 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { numb_models = models.size(); dfparam = DP_DeepSpinModelDeviGetDimFParam(dp); daparam = DP_DeepSpinModelDeviGetDimAParam(dp); + dchgspin = DP_DeepSpinModelDeviGetDimChgSpin(dp); aparam_nall = DP_DeepSpinModelDeviIsAParamNAll(dp); has_default_fparam_ = DP_DeepSpinModelDeviHasDefaultFParam(dp); dpbase = (DP_DeepBaseModelDevi*)dp; }; + /** + * @brief Get the dimension of the charge/spin input. + * @return The dimension of the charge/spin input (0 if the models have no + *charge/spin embedding). + **/ + int dim_chg_spin() const { + assert(dp); + return dchgspin; + } + /** * @brief Evaluate the energy, force, magnetic force and virial by using this *DP spin model deviation. @@ -2909,6 +3042,11 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { * nframes x natoms x dim_aparam. * natoms x dim_aparam. Then all frames are assumed to be provided with the *same aparam. + * @param[in] charge_spin The charge/spin input. The array can be of size: + * nframes x dim_chg_spin. + * dim_chg_spin. Then all frames are assumed to be provided with the same + *charge_spin. Leave it empty to use the model's stored + *default_chg_spin. **/ template void compute( @@ -2921,7 +3059,8 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()) { + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); @@ -2952,9 +3091,14 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + // charge_spin routes to the version-3 C API; nullptr keeps version-2 so + // non-charge_spin models still work against an older libdeepmd_c. + std::vector charge_spin_tiled_; + const VALUETYPE* charge_spin__ = validate_charge_spin( + charge_spin, dchgspin, nframes, charge_spin_tiled_); _DP_DeepSpinModelDeviCompute( - dp, natoms, coord_, spin_, atype_, box_, fparam__, aparam__, ener_, - force_, force_mag_, virial_, nullptr, nullptr); + dp, natoms, coord_, spin_, atype_, box_, fparam__, aparam__, + charge_spin__, ener_, force_, force_mag_, virial_, nullptr, nullptr); DP_CHECK_OK(DP_DeepSpinModelDeviCheckOK, dp); // reshape @@ -3002,6 +3146,11 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { * nframes x natoms x dim_aparam. * natoms x dim_aparam. Then all frames are assumed to be provided with the *same aparam. + * @param[in] charge_spin The charge/spin input. The array can be of size: + * nframes x dim_chg_spin. + * dim_chg_spin. Then all frames are assumed to be provided with the same + *charge_spin. Leave it empty to use the model's stored + *default_chg_spin. **/ template void compute( @@ -3016,7 +3165,8 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()) { + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); @@ -3051,9 +3201,15 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + // charge_spin routes to the version-3 C API; nullptr keeps version-2 so + // non-charge_spin models still work against an older libdeepmd_c. + std::vector charge_spin_tiled_; + const VALUETYPE* charge_spin__ = validate_charge_spin( + charge_spin, dchgspin, nframes, charge_spin_tiled_); _DP_DeepSpinModelDeviCompute( - dp, natoms, coord_, spin_, atype_, box_, fparam__, aparam__, ener_, - force_, force_mag_, virial_, atomic_ener_, atomic_virial_); + dp, natoms, coord_, spin_, atype_, box_, fparam__, aparam__, + charge_spin__, ener_, force_, force_mag_, virial_, atomic_ener_, + atomic_virial_); DP_CHECK_OK(DP_DeepSpinModelDeviCheckOK, dp); // reshape @@ -3113,6 +3269,11 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { * nframes x natoms x dim_aparam. * natoms x dim_aparam. Then all frames are assumed to be provided with the *same aparam. + * @param[in] charge_spin The charge/spin input. The array can be of size: + * nframes x dim_chg_spin. + * dim_chg_spin. Then all frames are assumed to be provided with the same + *charge_spin. Leave it empty to use the model's stored + *default_chg_spin. **/ template void compute( @@ -3128,7 +3289,8 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { const InputNlist& lmp_list, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()) { + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); @@ -3159,10 +3321,15 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + // charge_spin routes to the version-3 C API; nullptr keeps version-2 so + // non-charge_spin models still work against an older libdeepmd_c. + std::vector charge_spin_tiled_; + const VALUETYPE* charge_spin__ = validate_charge_spin( + charge_spin, dchgspin, nframes, charge_spin_tiled_); _DP_DeepSpinModelDeviComputeNList( dp, natoms, coord_, spin_, atype_, box_, nghost, lmp_list.nl, ago, - fparam__, aparam__, ener_, force_, force_mag_, virial_, nullptr, - nullptr); + fparam__, aparam__, charge_spin__, ener_, force_, force_mag_, virial_, + nullptr, nullptr); DP_CHECK_OK(DP_DeepSpinModelDeviCheckOK, dp); // reshape ener.resize(numb_models); @@ -3213,6 +3380,11 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { * nframes x natoms x dim_aparam. * natoms x dim_aparam. Then all frames are assumed to be provided with the *same aparam. + * @param[in] charge_spin The charge/spin input. The array can be of size: + * nframes x dim_chg_spin. + * dim_chg_spin. Then all frames are assumed to be provided with the same + *charge_spin. Leave it empty to use the model's stored + *default_chg_spin. **/ template void compute( @@ -3230,7 +3402,8 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { const InputNlist& lmp_list, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()) { + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); @@ -3266,10 +3439,15 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + // charge_spin routes to the version-3 C API; nullptr keeps version-2 so + // non-charge_spin models still work against an older libdeepmd_c. + std::vector charge_spin_tiled_; + const VALUETYPE* charge_spin__ = validate_charge_spin( + charge_spin, dchgspin, nframes, charge_spin_tiled_); _DP_DeepSpinModelDeviComputeNList( dp, natoms, coord_, spin_, atype_, box_, nghost, lmp_list.nl, ago, - fparam__, aparam__, ener_, force_, force_mag_, virial_, atomic_ener_, - atomic_virial_); + fparam__, aparam__, charge_spin__, ener_, force_, force_mag_, virial_, + atomic_ener_, atomic_virial_); DP_CHECK_OK(DP_DeepSpinModelDeviCheckOK, dp); // reshape ener.resize(numb_models); @@ -3305,6 +3483,7 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { private: DP_DeepSpinModelDevi* dp; + int dchgspin; }; /** diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index 1786099110..4e91591ad3 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -350,7 +350,8 @@ inline void DP_DeepSpinCompute_variant(DP_DeepSpin* dp, VALUETYPE* force_mag, VALUETYPE* virial, VALUETYPE* atomic_energy, - VALUETYPE* atomic_virial) { + VALUETYPE* atomic_virial, + const VALUETYPE* charge_spin = nullptr) { // init C++ vectors from C arrays std::vector coord_(coord, coord + nframes * natoms * 3); std::vector spin_(spin, spin + nframes * natoms * 3); @@ -368,11 +369,18 @@ inline void DP_DeepSpinCompute_variant(DP_DeepSpin* dp, if (aparam) { aparam_.assign(aparam, aparam + nframes * natoms * dp->daparam); } + // charge_spin is converted to double for the api_cc::DeepSpin interface + // (the compute backend stores/uses charge_spin as float64). + std::vector charge_spin_; + if (charge_spin) { + charge_spin_.assign(charge_spin, + charge_spin + nframes * dp->dp.dim_chg_spin()); + } std::vector e; std::vector f, fm, v, ae, av; DP_REQUIRES_OK(dp, dp->dp.compute(e, f, fm, v, ae, av, coord_, spin_, atype_, - cell_, fparam_, aparam_)); + cell_, fparam_, aparam_, charge_spin_)); // copy from C++ vectors to C arrays, if not NULL pointer if (energy) { std::copy(e.begin(), e.end(), energy); @@ -408,7 +416,8 @@ template void DP_DeepSpinCompute_variant(DP_DeepSpin* dp, double* force_mag, double* virial, double* atomic_energy, - double* atomic_virial); + double* atomic_virial, + const double* charge_spin); template void DP_DeepSpinCompute_variant(DP_DeepSpin* dp, const int nframes, @@ -424,7 +433,8 @@ template void DP_DeepSpinCompute_variant(DP_DeepSpin* dp, float* force_mag, float* virial, float* atomic_energy, - float* atomic_virial); + float* atomic_virial, + const float* charge_spin); template inline void DP_DeepPotComputeNList_variant( @@ -539,24 +549,26 @@ template void DP_DeepPotComputeNList_variant(DP_DeepPot* dp, // support spin template -inline void DP_DeepSpinComputeNList_variant(DP_DeepSpin* dp, - const int nframes, - const int natoms, - const VALUETYPE* coord, - const VALUETYPE* spin, - const int* atype, - const VALUETYPE* cell, - const int nghost, - const DP_Nlist* nlist, - const int ago, - const VALUETYPE* fparam, - const VALUETYPE* aparam, - double* energy, - VALUETYPE* force, - VALUETYPE* force_mag, - VALUETYPE* virial, - VALUETYPE* atomic_energy, - VALUETYPE* atomic_virial) { +inline void DP_DeepSpinComputeNList_variant( + DP_DeepSpin* dp, + const int nframes, + const int natoms, + const VALUETYPE* coord, + const VALUETYPE* spin, + const int* atype, + const VALUETYPE* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const VALUETYPE* fparam, + const VALUETYPE* aparam, + double* energy, + VALUETYPE* force, + VALUETYPE* force_mag, + VALUETYPE* virial, + VALUETYPE* atomic_energy, + VALUETYPE* atomic_virial, + const VALUETYPE* charge_spin = nullptr) { // init C++ vectors from C arrays std::vector coord_(coord, coord + nframes * natoms * 3); std::vector spin_(spin, spin + nframes * natoms * 3); @@ -577,11 +589,18 @@ inline void DP_DeepSpinComputeNList_variant(DP_DeepSpin* dp, (dp->aparam_nall ? natoms : (natoms - nghost)) * dp->daparam); } + // charge_spin is converted to double for the api_cc::DeepSpin interface + // (the compute backend stores/uses charge_spin as float64). + std::vector charge_spin_; + if (charge_spin) { + charge_spin_.assign(charge_spin, + charge_spin + nframes * dp->dp.dim_chg_spin()); + } std::vector e; std::vector f, fm, v, ae, av; - DP_REQUIRES_OK( - dp, dp->dp.compute(e, f, fm, v, ae, av, coord_, spin_, atype_, cell_, - nghost, nlist->nl, ago, fparam_, aparam_)); + DP_REQUIRES_OK(dp, dp->dp.compute(e, f, fm, v, ae, av, coord_, spin_, atype_, + cell_, nghost, nlist->nl, ago, fparam_, + aparam_, charge_spin_)); // copy from C++ vectors to C arrays, if not NULL pointer if (energy) { std::copy(e.begin(), e.end(), energy); @@ -602,24 +621,26 @@ inline void DP_DeepSpinComputeNList_variant(DP_DeepSpin* dp, std::copy(av.begin(), av.end(), atomic_virial); } } -template void DP_DeepSpinComputeNList_variant(DP_DeepSpin* dp, - const int nframes, - const int natoms, - const double* coord, - const double* spin, - const int* atype, - const double* cell, - const int nghost, - const DP_Nlist* nlist, - const int ago, - const double* fparam, - const double* aparam, - double* energy, - double* force, - double* force_mag, - double* virial, - double* atomic_energy, - double* atomic_virial); +template void DP_DeepSpinComputeNList_variant( + DP_DeepSpin* dp, + const int nframes, + const int natoms, + const double* coord, + const double* spin, + const int* atype, + const double* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const double* fparam, + const double* aparam, + double* energy, + double* force, + double* force_mag, + double* virial, + double* atomic_energy, + double* atomic_virial, + const double* charge_spin); template void DP_DeepSpinComputeNList_variant(DP_DeepSpin* dp, const int nframes, const int natoms, @@ -637,7 +658,8 @@ template void DP_DeepSpinComputeNList_variant(DP_DeepSpin* dp, float* force_mag, float* virial, float* atomic_energy, - float* atomic_virial); + float* atomic_virial, + const float* charge_spin); template inline void DP_DeepPotComputeMixedType_variant(DP_DeepPot* dp, @@ -863,21 +885,23 @@ template void DP_DeepPotModelDeviCompute_variant( const float* charge_spin); template -void DP_DeepSpinModelDeviCompute_variant(DP_DeepSpinModelDevi* dp, - const int nframes, - const int natoms, - const VALUETYPE* coord, - const VALUETYPE* spin, - const int* atype, - const VALUETYPE* cell, - const VALUETYPE* fparam, - const VALUETYPE* aparam, - double* energy, - VALUETYPE* force, - VALUETYPE* force_mag, - VALUETYPE* virial, - VALUETYPE* atomic_energy, - VALUETYPE* atomic_virial) { +void DP_DeepSpinModelDeviCompute_variant( + DP_DeepSpinModelDevi* dp, + const int nframes, + const int natoms, + const VALUETYPE* coord, + const VALUETYPE* spin, + const int* atype, + const VALUETYPE* cell, + const VALUETYPE* fparam, + const VALUETYPE* aparam, + double* energy, + VALUETYPE* force, + VALUETYPE* force_mag, + VALUETYPE* virial, + VALUETYPE* atomic_energy, + VALUETYPE* atomic_virial, + const VALUETYPE* charge_spin = nullptr) { if (!validate_model_devi_nframes(dp, nframes)) { return; } @@ -898,16 +922,24 @@ void DP_DeepSpinModelDeviCompute_variant(DP_DeepSpinModelDevi* dp, if (aparam) { aparam_.assign(aparam, aparam + nframes * natoms * dp->daparam); } + // charge_spin is converted to double for the api_cc::DeepSpinModelDevi + // interface (the compute backend stores/uses charge_spin as float64). + std::vector charge_spin_; + if (charge_spin) { + charge_spin_.assign(charge_spin, + charge_spin + nframes * dp->dp.dim_chg_spin()); + } // different from DeepPot std::vector e; std::vector> f, fm, v, ae, av; if (atomic_energy || atomic_virial) { - DP_REQUIRES_OK(dp, dp->dp.compute(e, f, fm, v, ae, av, coord_, spin_, - atype_, cell_, fparam_, aparam_)); + DP_REQUIRES_OK( + dp, dp->dp.compute(e, f, fm, v, ae, av, coord_, spin_, atype_, cell_, + fparam_, aparam_, charge_spin_)); } else { DP_REQUIRES_OK(dp, dp->dp.compute(e, f, fm, v, coord_, spin_, atype_, cell_, - fparam_, aparam_)); + fparam_, aparam_, charge_spin_)); } // 2D vector to 2D array, flatten first if (energy) { @@ -955,7 +987,8 @@ template void DP_DeepSpinModelDeviCompute_variant( double* force_mag, double* virial, double* atomic_energy, - double* atomic_virial); + double* atomic_virial, + const double* charge_spin); template void DP_DeepSpinModelDeviCompute_variant( DP_DeepSpinModelDevi* dp, @@ -972,7 +1005,8 @@ template void DP_DeepSpinModelDeviCompute_variant( float* force_mag, float* virial, float* atomic_energy, - float* atomic_virial); + float* atomic_virial, + const float* charge_spin); template void DP_DeepPotModelDeviComputeNList_variant( @@ -1100,24 +1134,26 @@ template void DP_DeepPotModelDeviComputeNList_variant( // support spin multi model. template -void DP_DeepSpinModelDeviComputeNList_variant(DP_DeepSpinModelDevi* dp, - const int nframes, - const int natoms, - const VALUETYPE* coord, - const VALUETYPE* spin, - const int* atype, - const VALUETYPE* cell, - const int nghost, - const DP_Nlist* nlist, - const int ago, - const VALUETYPE* fparam, - const VALUETYPE* aparam, - double* energy, - VALUETYPE* force, - VALUETYPE* force_mag, - VALUETYPE* virial, - VALUETYPE* atomic_energy, - VALUETYPE* atomic_virial) { +void DP_DeepSpinModelDeviComputeNList_variant( + DP_DeepSpinModelDevi* dp, + const int nframes, + const int natoms, + const VALUETYPE* coord, + const VALUETYPE* spin, + const int* atype, + const VALUETYPE* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const VALUETYPE* fparam, + const VALUETYPE* aparam, + double* energy, + VALUETYPE* force, + VALUETYPE* force_mag, + VALUETYPE* virial, + VALUETYPE* atomic_energy, + VALUETYPE* atomic_virial, + const VALUETYPE* charge_spin = nullptr) { if (!validate_model_devi_nframes(dp, nframes)) { return; } @@ -1140,17 +1176,24 @@ void DP_DeepSpinModelDeviComputeNList_variant(DP_DeepSpinModelDevi* dp, aparam, aparam + (dp->aparam_nall ? natoms : (natoms - nghost)) * dp->daparam); } + // charge_spin is converted to double for the api_cc::DeepSpinModelDevi + // interface (the compute backend stores/uses charge_spin as float64). + std::vector charge_spin_; + if (charge_spin) { + charge_spin_.assign(charge_spin, + charge_spin + nframes * dp->dp.dim_chg_spin()); + } // different from DeepPot std::vector e; std::vector> f, fm, v, ae, av; if (atomic_energy || atomic_virial) { - DP_REQUIRES_OK( - dp, dp->dp.compute(e, f, fm, v, ae, av, coord_, spin_, atype_, cell_, - nghost, nlist->nl, ago, fparam_, aparam_)); + DP_REQUIRES_OK(dp, dp->dp.compute(e, f, fm, v, ae, av, coord_, spin_, + atype_, cell_, nghost, nlist->nl, ago, + fparam_, aparam_, charge_spin_)); } else { DP_REQUIRES_OK( dp, dp->dp.compute(e, f, fm, v, coord_, spin_, atype_, cell_, nghost, - nlist->nl, ago, fparam_, aparam_)); + nlist->nl, ago, fparam_, aparam_, charge_spin_)); } // 2D vector to 2D array, flatten first if (energy) { @@ -1200,7 +1243,8 @@ template void DP_DeepSpinModelDeviComputeNList_variant( double* force_mag, double* virial, double* atomic_energy, - double* atomic_virial); + double* atomic_virial, + const double* charge_spin); template void DP_DeepSpinModelDeviComputeNList_variant( DP_DeepSpinModelDevi* dp, const int nframes, @@ -1219,7 +1263,8 @@ template void DP_DeepSpinModelDeviComputeNList_variant( float* force_mag, float* virial, float* atomic_energy, - float* atomic_virial); + float* atomic_virial, + const float* charge_spin); template inline void DP_DeepTensorComputeTensor_variant(DP_DeepTensor* dt, @@ -2034,6 +2079,100 @@ void DP_DeepSpinComputeNListf2(DP_DeepSpin* dp, aparam, energy, force, force_mag, virial, atomic_energy, atomic_virial); } +// charge_spin-aware spin variants (version 3): same as the version-2 spin +// functions plus a per-frame charge_spin input (nframes x dim_chg_spin). +void DP_DeepSpinCompute3(DP_DeepSpin* dp, + const int nframes, + const int natoms, + const double* coord, + const double* spin, + const int* atype, + const double* cell, + const double* fparam, + const double* aparam, + const double* charge_spin, + double* energy, + double* force, + double* force_mag, + double* virial, + double* atomic_energy, + double* atomic_virial) { + DP_DeepSpinCompute_variant( + dp, nframes, natoms, coord, spin, atype, cell, fparam, aparam, energy, + force, force_mag, virial, atomic_energy, atomic_virial, charge_spin); +} + +void DP_DeepSpinComputef3(DP_DeepSpin* dp, + const int nframes, + const int natoms, + const float* coord, + const float* spin, + const int* atype, + const float* cell, + const float* fparam, + const float* aparam, + const float* charge_spin, + double* energy, + float* force, + float* force_mag, + float* virial, + float* atomic_energy, + float* atomic_virial) { + DP_DeepSpinCompute_variant( + dp, nframes, natoms, coord, spin, atype, cell, fparam, aparam, energy, + force, force_mag, virial, atomic_energy, atomic_virial, charge_spin); +} + +void DP_DeepSpinComputeNList3(DP_DeepSpin* dp, + const int nframes, + const int natoms, + const double* coord, + const double* spin, + const int* atype, + const double* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const double* fparam, + const double* aparam, + const double* charge_spin, + double* energy, + double* force, + double* force_mag, + double* virial, + double* atomic_energy, + double* atomic_virial) { + DP_DeepSpinComputeNList_variant( + dp, nframes, natoms, coord, spin, atype, cell, nghost, nlist, ago, fparam, + aparam, energy, force, force_mag, virial, atomic_energy, atomic_virial, + charge_spin); +} + +void DP_DeepSpinComputeNListf3(DP_DeepSpin* dp, + const int nframes, + const int natoms, + const float* coord, + const float* spin, + const int* atype, + const float* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const float* fparam, + const float* aparam, + const float* charge_spin, + double* energy, + float* force, + float* force_mag, + float* virial, + float* atomic_energy, + float* atomic_virial) { + DP_DeepSpinComputeNList_variant( + dp, nframes, natoms, coord, spin, atype, cell, nghost, nlist, ago, fparam, + aparam, energy, force, force_mag, virial, atomic_energy, atomic_virial, + charge_spin); +} + // end multiple frames void DP_DeepPotComputeMixedType(DP_DeepPot* dp, @@ -2385,6 +2524,101 @@ void DP_DeepSpinModelDeviComputeNListf2(DP_DeepSpinModelDevi* dp, aparam, energy, force, force_mag, virial, atomic_energy, atomic_virial); } +// charge_spin-aware spin model deviation variants (version 3): same as the +// version-2 functions plus a per-frame charge_spin input (nframes x +// dim_chg_spin). +void DP_DeepSpinModelDeviCompute3(DP_DeepSpinModelDevi* dp, + const int nframes, + const int natoms, + const double* coord, + const double* spin, + const int* atype, + const double* cell, + const double* fparam, + const double* aparam, + const double* charge_spin, + double* energy, + double* force, + double* force_mag, + double* virial, + double* atomic_energy, + double* atomic_virial) { + DP_DeepSpinModelDeviCompute_variant( + dp, nframes, natoms, coord, spin, atype, cell, fparam, aparam, energy, + force, force_mag, virial, atomic_energy, atomic_virial, charge_spin); +} + +void DP_DeepSpinModelDeviComputef3(DP_DeepSpinModelDevi* dp, + const int nframes, + const int natoms, + const float* coord, + const float* spin, + const int* atype, + const float* cell, + const float* fparam, + const float* aparam, + const float* charge_spin, + double* energy, + float* force, + float* force_mag, + float* virial, + float* atomic_energy, + float* atomic_virial) { + DP_DeepSpinModelDeviCompute_variant( + dp, nframes, natoms, coord, spin, atype, cell, fparam, aparam, energy, + force, force_mag, virial, atomic_energy, atomic_virial, charge_spin); +} + +void DP_DeepSpinModelDeviComputeNList3(DP_DeepSpinModelDevi* dp, + const int nframes, + const int natoms, + const double* coord, + const double* spin, + const int* atype, + const double* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const double* fparam, + const double* aparam, + const double* charge_spin, + double* energy, + double* force, + double* force_mag, + double* virial, + double* atomic_energy, + double* atomic_virial) { + DP_DeepSpinModelDeviComputeNList_variant( + dp, nframes, natoms, coord, spin, atype, cell, nghost, nlist, ago, fparam, + aparam, energy, force, force_mag, virial, atomic_energy, atomic_virial, + charge_spin); +} + +void DP_DeepSpinModelDeviComputeNListf3(DP_DeepSpinModelDevi* dp, + const int nframes, + const int natoms, + const float* coord, + const float* spin, + const int* atype, + const float* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const float* fparam, + const float* aparam, + const float* charge_spin, + double* energy, + float* force, + float* force_mag, + float* virial, + float* atomic_energy, + float* atomic_virial) { + DP_DeepSpinModelDeviComputeNList_variant( + dp, nframes, natoms, coord, spin, atype, cell, nghost, nlist, ago, fparam, + aparam, energy, force, force_mag, virial, atomic_energy, atomic_virial, + charge_spin); +} + // base model methods const char* DP_DeepBaseModelGetTypeMap(DP_DeepBaseModel* dpbase) { std::string type_map; @@ -2483,6 +2717,8 @@ int DP_DeepPotGetDimAParam(DP_DeepPot* dp) { int DP_DeepPotGetDimChgSpin(DP_DeepPot* dp) { return dp->dp.dim_chg_spin(); } +int DP_DeepSpinGetDimChgSpin(DP_DeepSpin* dp) { return dp->dp.dim_chg_spin(); } + bool DP_DeepPotIsAParamNAll(DP_DeepPot* dp) { return DP_DeepBaseModelIsAParamNAll(static_cast(dp)); } @@ -2598,6 +2834,10 @@ int DP_DeepSpinModelDeviGetDimAParam(DP_DeepSpinModelDevi* dp) { static_cast(dp)); } +int DP_DeepSpinModelDeviGetDimChgSpin(DP_DeepSpinModelDevi* dp) { + return dp->dp.dim_chg_spin(); +} + bool DP_DeepSpinModelDeviIsAParamNAll(DP_DeepSpinModelDevi* dp) { return DP_DeepBaseModelDeviIsAParamNAll( static_cast(dp)); diff --git a/source/api_cc/include/DeepSpin.h b/source/api_cc/include/DeepSpin.h index a4f38461ca..2a13b79cf8 100644 --- a/source/api_cc/include/DeepSpin.h +++ b/source/api_cc/include/DeepSpin.h @@ -161,6 +161,91 @@ class DeepSpinBackend : public DeepBaseModelBackend { const bool atomic) = 0; /** @} */ + /** + * @brief Get dimension of charge/spin condition inputs. + * Returns 0 for backends that do not support charge/spin conditioning. + **/ + virtual int dim_chg_spin() const { return 0; } + + // charge_spin-aware computew overloads. Default implementations call the + // existing pure-virtual overloads (ignoring charge_spin) so that backends + // that do not support charge/spin do not need any changes. DeepSpinPTExpt + // overrides these to thread charge_spin through to the model. + virtual void computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + computew(ener, force, force_mag, virial, atom_energy, atom_virial, coord, + spin, atype, box, fparam, aparam, atomic); + } + virtual void computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + computew(ener, force, force_mag, virial, atom_energy, atom_virial, coord, + spin, atype, box, fparam, aparam, atomic); + } + virtual void computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + computew(ener, force, force_mag, virial, atom_energy, atom_virial, coord, + spin, atype, box, nghost, inlist, ago, fparam, aparam, atomic); + } + virtual void computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + computew(ener, force, force_mag, virial, atom_energy, atom_virial, coord, + spin, atype, box, nghost, inlist, ago, fparam, aparam, atomic); + } + /** * @brief Get the per-type use_spin flags. * @return A vector of booleans indicating which atom types have spin enabled. @@ -222,6 +307,10 @@ class DeepSpin : public DeepBaseModel { * nframes x natoms x dim_aparam. * natoms x dim_aparam. Then all frames are assumed to be provided with the *same aparam. + * @param[in] charge_spin The charge/spin parameter. The array can be of size + *nframes x dim_chg_spin. + * dim_chg_spin. Then all frames are assumed to be provided with the same + *charge_spin. Leave it empty to use the model's stored default_chg_spin. * @{ **/ template @@ -234,7 +323,8 @@ class DeepSpin : public DeepBaseModel { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); template void compute(std::vector& ener, std::vector& force, @@ -245,7 +335,8 @@ class DeepSpin : public DeepBaseModel { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); /** @} */ /** @@ -273,6 +364,10 @@ class DeepSpin : public DeepBaseModel { * nframes x natoms x dim_aparam. * natoms x dim_aparam. Then all frames are assumed to be provided with the *same aparam. + * @param[in] charge_spin The charge/spin parameter. The array can be of size + *nframes x dim_chg_spin. + * dim_chg_spin. Then all frames are assumed to be provided with the same + *charge_spin. Leave it empty to use the model's stored default_chg_spin. * @{ **/ template @@ -288,7 +383,8 @@ class DeepSpin : public DeepBaseModel { const InputNlist& inlist, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); template void compute(std::vector& ener, std::vector& force, @@ -302,7 +398,8 @@ class DeepSpin : public DeepBaseModel { const InputNlist& inlist, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); /** @} */ /** @@ -329,6 +426,10 @@ class DeepSpin : public DeepBaseModel { * nframes x natoms x dim_aparam. * natoms x dim_aparam. Then all frames are assumed to be provided with the *same aparam. + * @param[in] charge_spin The charge/spin parameter. The array can be of size + *nframes x dim_chg_spin. + * dim_chg_spin. Then all frames are assumed to be provided with the same + *charge_spin. Leave it empty to use the model's stored default_chg_spin. * @{ **/ template @@ -343,7 +444,8 @@ class DeepSpin : public DeepBaseModel { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); template void compute(std::vector& ener, std::vector& force, @@ -356,7 +458,8 @@ class DeepSpin : public DeepBaseModel { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); /** @} */ /** @@ -386,6 +489,10 @@ class DeepSpin : public DeepBaseModel { * nframes x natoms x dim_aparam. * natoms x dim_aparam. Then all frames are assumed to be provided with the *same aparam. + * @param[in] charge_spin The charge/spin parameter. The array can be of size + *nframes x dim_chg_spin. + * dim_chg_spin. Then all frames are assumed to be provided with the same + *charge_spin. Leave it empty to use the model's stored default_chg_spin. * @{ **/ template @@ -403,7 +510,8 @@ class DeepSpin : public DeepBaseModel { const InputNlist& lmp_list, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); template void compute(std::vector& ener, std::vector& force, @@ -419,9 +527,17 @@ class DeepSpin : public DeepBaseModel { const InputNlist& lmp_list, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); /** @} */ + /** + * @brief Get dimension of the charge/spin condition inputs. + * @return The dimension of charge_spin; 0 when the model does not support + *charge/spin conditioning. + **/ + int dim_chg_spin() const; + /** * @brief Get the per-type use_spin flags. * @return A vector of booleans indicating which atom types have spin enabled. @@ -461,6 +577,17 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { const int& gpu_rank = 0, const std::vector& file_contents = std::vector()); + + /** + * @brief Get the dimension of the charge/spin input. + * @return The dimension of the charge/spin input (0 if the models have no + *charge/spin embedding). Taken from the first model; all models are assumed + *to share the same value. + **/ + int dim_chg_spin() const { + return numb_models > 0 ? dps[0]->dim_chg_spin() : 0; + }; + /** * @brief Evaluate the energy, force and virial by using these DP spin models. * @param[out] all_ener The system energies of all models. @@ -483,6 +610,10 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { * natoms x dim_aparam. Then all frames are assumed to be provided with the *same aparam. dim_aparam. Then all frames and atoms are provided with the *same aparam. + * @param[in] charge_spin The charge/spin parameter. The array can be of size + *nframes x dim_chg_spin. + * dim_chg_spin. Then all frames are assumed to be provided with the same + *charge_spin. Leave it empty to use the model's stored default_chg_spin. **/ template void compute(std::vector& all_ener, @@ -494,7 +625,8 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); /** * @brief Evaluate the energy, force, virial, atomic energy, and atomic virial @@ -521,6 +653,10 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { * natoms x dim_aparam. Then all frames are assumed to be provided with the *same aparam. dim_aparam. Then all frames and atoms are provided with the *same aparam. + * @param[in] charge_spin The charge/spin parameter. The array can be of size + *nframes x dim_chg_spin. + * dim_chg_spin. Then all frames are assumed to be provided with the same + *charge_spin. Leave it empty to use the model's stored default_chg_spin. **/ template void compute(std::vector& all_ener, @@ -534,7 +670,8 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); /** * @brief Evaluate the energy, force, magnetic force and virial by using these @@ -562,6 +699,10 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { * natoms x dim_aparam. Then all frames are assumed to be provided with the *same aparam. dim_aparam. Then all frames and atoms are provided with the *same aparam. + * @param[in] charge_spin The charge/spin parameter. The array can be of size + *nframes x dim_chg_spin. + * dim_chg_spin. Then all frames are assumed to be provided with the same + *charge_spin. Leave it empty to use the model's stored default_chg_spin. **/ template void compute(std::vector& all_ener, @@ -576,7 +717,8 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { const InputNlist& lmp_list, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); /** * @brief Evaluate the energy, force, magnetic force, virial, atomic energy, @@ -606,6 +748,10 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { * natoms x dim_aparam. Then all frames are assumed to be provided with the *same aparam. dim_aparam. Then all frames and atoms are provided with the *same aparam. + * @param[in] charge_spin The charge/spin parameter. The array can be of size + *nframes x dim_chg_spin. + * dim_chg_spin. Then all frames are assumed to be provided with the same + *charge_spin. Leave it empty to use the model's stored default_chg_spin. **/ template void compute(std::vector& all_ener, @@ -622,7 +768,8 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { const InputNlist& lmp_list, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); /** * @brief Get the per-type use_spin flags from the first model. diff --git a/source/api_cc/include/DeepSpinPTExpt.h b/source/api_cc/include/DeepSpinPTExpt.h index a4d0311d98..500c492607 100644 --- a/source/api_cc/include/DeepSpinPTExpt.h +++ b/source/api_cc/include/DeepSpinPTExpt.h @@ -59,6 +59,7 @@ class DeepSpinPTExpt : public DeepSpinBackend { const int& ago, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); /** * @brief Evaluate without nlist (standalone — builds nlist, folds back). @@ -76,6 +77,7 @@ class DeepSpinPTExpt : public DeepSpinBackend { const std::vector& box, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); public: @@ -99,6 +101,10 @@ class DeepSpinPTExpt : public DeepSpinBackend { assert(inited); return daparam; }; + int dim_chg_spin() const override { + assert(inited); + return dchgspin; + }; void get_type_map(std::string& type_map); bool is_aparam_nall() const { assert(inited); @@ -113,7 +119,8 @@ class DeepSpinPTExpt : public DeepSpinBackend { return has_default_fparam_; }; - // forward to template class + // forward to template class (no charge_spin — uses default_chg_spin_ + // fallback) void computew(std::vector& ener, std::vector& force, std::vector& force_mag, @@ -173,13 +180,77 @@ class DeepSpinPTExpt : public DeepSpinBackend { const std::vector& aparam, const bool atomic); + // charge_spin overloads — pass runtime charge/spin per call + void computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) override; + void computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) override; + void computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) override; + void computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) override; + private: bool inited; int ntypes; int ntypes_spin; int dfparam; int daparam; - int dim_chg_spin; + int dchgspin; bool aparam_nall; bool has_default_fparam_; std::vector default_fparam_; @@ -194,6 +265,13 @@ class DeepSpinPTExpt : public DeepSpinBackend { // Whether the exported graph consumes the compact edge schema (native spin, // shared with DeepPotPTExpt) rather than the deepspin-scheme nlist contract. bool lower_input_is_edge_ = false; + // Whether the exported graph consumes the NeighborGraph schema (native + // spin; single-rank runs run_model_graph, multi-rank runs + // run_model_graph_with_comm). Mutually exclusive with lower_input_is_edge_. + bool lower_input_is_graph_ = false; + // Edge-vector precision recorded by the .pt2 artifact for the graph route + // (mirrors DeepPotPTExpt); read from metadata ``graph_edge_dtype``. + bool graph_edge_fp32_ = false; int nnei; // expected nlist nnei dimension (= sum(sel)) NeighborListData nlist_data; at::Tensor mapping_tensor; // cached mapping tensor (LAMMPS path) @@ -206,6 +284,22 @@ class DeepSpinPTExpt : public DeepSpinBackend { // Mirrors descriptor's has_message_passing(). See DeepPotPTExpt.h // for the full rationale and gating role. bool has_message_passing_ = false; + // Whether the collective empty-rank preflight (allreduce of the minimum + // owned+ghost node count over the LAMMPS communicator, native-spin graph + // with-comm route) has PASSED for the current neighbor topology. Twin of + // ``DeepPotPTExpt::graph_comm_preflight_done_``: the preflight re-runs + // whenever the with-comm graph branch is entered with ``ago == 0`` (every + // topology rebuild) or before its first success, because the node count + // shares the lifetime of the cached nlist/mapping/edge topology and + // re-running the collective on cache-hit (``ago > 0``) force calls would + // add a global synchronization per MD step with no added protection. + bool graph_comm_preflight_done_ = false; + // Model-level pair-type exclusion keep table, built ONCE in ``init`` from + // the ``pair_exclude_types`` metadata field (see DeepPotPTExpt.h for the + // full rationale). UNDEFINED => no exclusion (identity). Applied at the + // neighbor-graph BUILD seam (``applyPairExclusion``), never inside the + // exported lower. + torch::Tensor pair_exclude_table_; std::unique_ptr with_comm_tempfile_; std::unique_ptr with_comm_loader; @@ -215,7 +309,8 @@ class DeepSpinPTExpt : public DeepSpinBackend { const torch::Tensor& nlist, const torch::Tensor& mapping, const torch::Tensor& fparam, - const torch::Tensor& aparam); + const torch::Tensor& aparam, + const torch::Tensor& charge_spin); /** * @brief Run the native-spin edge artifact: the energy edge schema plus the @@ -230,7 +325,60 @@ class DeepSpinPTExpt : public DeepSpinBackend { const torch::Tensor& edge_mask, const torch::Tensor& spin, const torch::Tensor& fparam, - const torch::Tensor& aparam); + const torch::Tensor& aparam, + const torch::Tensor& charge_spin); + + /** + * @brief Run the native-spin NeighborGraph artifact: the 10 base graph + * tensors (see commonPT.h GraphTensorPack), the per-node spin leaf + * (always present, positional index 10), then the conditional + * fparam/aparam/charge_spin tail -- combined native-spin + charge-spin + * FiLM models carry charge_spin at slot 13, mirroring the energy graph + * ABI's optional charge_spin tail. Single-rank; the multi-rank twin is + * ``run_model_graph_with_comm``. + */ + std::vector run_model_graph( + const torch::Tensor& atype, + const torch::Tensor& n_node, + const torch::Tensor& n_local, + const torch::Tensor& edge_index, + const torch::Tensor& edge_vec, + const torch::Tensor& edge_mask, + const torch::Tensor& destination_order, + const torch::Tensor& destination_row_ptr, + const torch::Tensor& source_order, + const torch::Tensor& source_row_ptr, + const torch::Tensor& spin, + const torch::Tensor& fparam, + const torch::Tensor& aparam, + const torch::Tensor& charge_spin); + + /** + * @brief Run the native-spin parallel GRAPH artifact: run_model_graph's + * ABI (spin at positional index 10) with the 8 border_op comm tensors + * appended after the conditional fparam/aparam/charge_spin tail. + * + * ``spin`` is the EXTENDED per-node spin -- ghost rows carry their + * owner's value from the LAMMPS ``sp`` forward-comm, so spin itself needs + * no cross-rank exchange; only the per-block ghost FEATURE refresh rides + * ``border_op``. Twin of ``DeepPotPTExpt::run_model_graph_with_comm``. + */ + std::vector run_model_graph_with_comm( + const torch::Tensor& atype, + const torch::Tensor& n_node, + const torch::Tensor& n_local, + const torch::Tensor& edge_index, + const torch::Tensor& edge_vec, + const torch::Tensor& edge_mask, + const torch::Tensor& destination_order, + const torch::Tensor& destination_row_ptr, + const torch::Tensor& source_order, + const torch::Tensor& source_row_ptr, + const torch::Tensor& spin, + const torch::Tensor& fparam, + const torch::Tensor& aparam, + const torch::Tensor& charge_spin, + const std::vector& comm_tensors); /** * @brief Run the native-spin parallel edge artifact: the energy edge @@ -248,10 +396,11 @@ class DeepSpinPTExpt : public DeepSpinBackend { const torch::Tensor& spin, const torch::Tensor& fparam, const torch::Tensor& aparam, + const torch::Tensor& charge_spin, const std::vector& comm_tensors); /** - * @brief Run with-comm spin artifact: 5-7 base inputs (incl. + * @brief Run with-comm spin artifact: 5-8 base inputs (incl. * extended_spin) + 8 comm tensors. */ std::vector run_model_with_comm( @@ -262,6 +411,7 @@ class DeepSpinPTExpt : public DeepSpinBackend { const torch::Tensor& mapping, const torch::Tensor& fparam, const torch::Tensor& aparam, + const torch::Tensor& charge_spin, const std::vector& comm_tensors); void extract_outputs(std::map& output_map, diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index 806caab494..066c6ba991 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -581,6 +582,94 @@ inline GraphTensorPack buildGraphTensors( return pack; } +/** + * @brief Normalize a graph-route aparam tensor to the flat node axis. + * + * The NeighborGraph ABI carries atomic parameters FLAT on the node axis -- + * shape (N, daparam) with N == n_node_count, the same axis as ``atype`` + * (mirrors ``build_synthetic_graph_inputs`` / ``_build_graph_dynamic_shapes`` + * on the Python export side). The runtime aparam carries the owned (local) + * rows only (``aparam_nall`` is structurally false for pt_expt models, see + * ``init``); on extended-region graphs (multi-rank routes, N == nall_real) + * the ghost rows are zero-padded here. Ghost fitting outputs are never + * retained -- the with-comm artifact masks non-owned energies before + * reduction, and the plain multi-rank remap sums energy over the owned + * prefix only -- so the padded values are inert. + * + * A ghost-only rank (``nlocal == 0``, ``N > 0``) synthesizes a full zero + * tensor: the graph still carries N nodes and the artifact requires the + * (N, daparam) input, while the owned-node mask keeps its contribution + * exactly zero. A missing aparam on a rank that OWNS atoms, or a width + * mismatch, is an explicit error -- the former silently returned the empty + * tensor (artifact reshape failure mid-collective), and a width mismatch + * used to be absorbed by broadcasting ``copy_`` (silent result corruption + * for daparam > 1). + * + * Shared by ``DeepPotPTExpt`` and ``DeepSpinPTExpt``: both graph routes carry + * aparam on the same flat node axis. + */ +inline at::Tensor extend_graph_aparam(const at::Tensor& aparam_tensor, + std::int64_t n_node_count, + std::int64_t nlocal, + std::int64_t daparam) { + if (daparam <= 0) { + return aparam_tensor; // model has no aparam input; passed through empty + } + if (aparam_tensor.numel() == 0) { + if (nlocal > 0) { + throw deepmd::deepmd_exception( + "aparam is required (dim_aparam=" + std::to_string(daparam) + + ") but no values were provided on a rank owning " + + std::to_string(nlocal) + " atoms."); + } + // ghost-only rank: there are no owned rows to supply; zeros are inert + // under the owned-node mask but the artifact needs the full node axis. + return torch::zeros({n_node_count, daparam}, aparam_tensor.options()); + } + if (aparam_tensor.numel() != nlocal * daparam) { + throw deepmd::deepmd_exception( + "aparam holds " + std::to_string(aparam_tensor.numel()) + + " values but the graph route expects nlocal * dim_aparam = " + + std::to_string(nlocal) + " * " + std::to_string(daparam) + "."); + } + at::Tensor owned = aparam_tensor.reshape({nlocal, daparam}); + if (nlocal == n_node_count) { + return owned; // single-rank / folded graph: nothing to pad + } + at::Tensor padded = + torch::zeros({n_node_count, daparam}, aparam_tensor.options()); + padded.slice(0, 0, nlocal).copy_(owned); + return padded; +} + +/** + * @brief Assert the flat graph-route aparam contract at the C++ boundary. + * + * The graph artifacts consume aparam FLAT on the node axis, shape + * (N, daparam) -- the layout ``extend_graph_aparam`` produces. A caller + * hand-rolling a rectangular (1, N, daparam) tensor (the pre-flat + * convention) would otherwise fail DEEP inside the artifact -- or, on a + * GPU-only route, only at deployment where no CPU test can catch it (the + * device-edge branch shipped exactly that bug). Failing loudly here turns + * any future such site into an immediate, self-explanatory error. + */ +inline void check_graph_aparam_flat(const at::Tensor& aparam, + std::int64_t daparam, + const char* where) { + if (daparam <= 0) { + return; + } + if (aparam.dim() != 2 || aparam.size(1) != daparam) { + std::ostringstream oss; + oss << where + << ": graph-route aparam must be flat (N, daparam) on the node axis " + "(produce it with extend_graph_aparam); got a rank-" + << aparam.dim() << " tensor of shape " << aparam.sizes() + << " for daparam = " << daparam << "."; + throw deepmd::deepmd_exception(oss.str()); + } +} + /** * @brief Graph pair-type exclusion: AND the per-edge keep-mask into * ``edge_mask``. @@ -754,6 +843,50 @@ inline void remap_graph_outputs_to_dense_keys( } } +/** + * @brief Remap NeighborGraph (graph-schema) native-spin public outputs onto + * the dense internal-key layout ``DeepSpinPTExpt::compute`` consumes. + * + * The single-rank native-spin graph forward is LOCAL-only and additionally + * emits ``force_mag`` (N, 3): per-node + * magnetic force, N == nloc, already exactly zero on non-spin-carrying atoms + * (the model's own type gate, not re-masked here per the project's + * one-owner design principle). + * + * Delegates the energy/force/virial/atom_virial remap to + * ``remap_graph_outputs_to_dense_keys`` and additionally zero-pads + * ``force_mag`` up to ``nall`` exactly like ``force`` (ghost rows already + * folded onto their local owners via ``edge_index``), writing it to + * ``energy_derv_r_mag`` -- the key ``DeepSpinPTExpt::compute`` reads for + * every other lower schema (dense nlist / edge_vec). + * + * **Single-rank only** (``fold_to_local=true``, so ``N == nloc``). The + * multi-rank sibling is + * ``remap_graph_spin_outputs_to_dense_keys_extended``; calling THIS one on an + * extended-region result throws on the ``index_put_`` below as soon as + * ``nloc < nall``, because ``force_mag_pub`` then carries ``nall`` rows. + * + * @param[in,out] output_map Output tensor map (public keys in, internal keys + * added). + * @param[in] nloc Number of local atoms (== N, the graph node count). + * @param[in] nall Extended atom count to pad the per-atom outputs up to. + * @param[in] atomic Whether atomic energy / virial were requested. + */ +inline void remap_graph_spin_outputs_to_dense_keys( + std::map& output_map, + const std::int64_t nloc, + const std::int64_t nall, + const bool atomic) { + using torch::indexing::Slice; + const std::int64_t nf = 1; + remap_graph_outputs_to_dense_keys(output_map, nloc, nall, atomic, + /*single_rank=*/true); + const auto& force_mag_pub = output_map.at("force_mag"); // (N, 3) + auto force_mag_full = torch::zeros({nf, nall, 1, 3}, force_mag_pub.options()); + force_mag_full.index_put_({0, Slice(0, nloc), 0}, force_mag_pub); + output_map["energy_derv_r_mag"] = force_mag_full; +} + /** * @brief Remap NeighborGraph public outputs onto the dense internal-key layout * for the MULTI-RANK (extended-region) path. @@ -816,6 +949,38 @@ inline void remap_graph_outputs_to_dense_keys_extended( } } +/** + * @brief Native-spin twin of ``remap_graph_outputs_to_dense_keys_extended``: + * the MULTI-RANK (extended-region) graph-spin output remap. + * + * Built with ``fold_to_local=false``, the graph has ``N == nall`` nodes, so + * ``force_mag`` is already the EXTENDED magnetic force -- one row per + * extended atom. Unlike the single-rank helper it must NOT zero-pad from + * ``nloc`` to ``nall``: the rows are already there, and padding would both + * truncate real ghost rows and mis-shape the tensor (the single-rank helper's + * ``index_put_({0, Slice(0, nloc), 0}, force_mag_pub)`` fails outright when + * ``force_mag_pub`` carries ``nall`` rows and ``nloc < nall``). + * + * Ghost magnetic-force rows stay distinct and are folded onto their owners by + * the LAMMPS spin reverse-comm, exactly as ghost conservative-force rows are. + * + * @param[in,out] output_map Output tensor map (public keys in, internal keys + * added). + * @param[in] nloc Number of local atoms (owned by this rank). + * @param[in] nall Extended atom count (== N, the graph node count). + * @param[in] atomic Whether atomic energy / virial were requested. + */ +inline void remap_graph_spin_outputs_to_dense_keys_extended( + std::map& output_map, + const std::int64_t nloc, + const std::int64_t nall, + const bool atomic) { + const std::int64_t nf = 1; + remap_graph_outputs_to_dense_keys_extended(output_map, nloc, nall, atomic); + const auto& force_mag_pub = output_map.at("force_mag"); // (N==nall, 3) + output_map["energy_derv_r_mag"] = force_mag_pub.reshape({nf, nall, 1, 3}); +} + } // namespace deepmd #endif // BUILD_PYTORCH diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index c1049cdfea..b2ad13336c 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -29,90 +29,9 @@ using deepmd::ptexpt::read_zip_entry; using namespace deepmd; namespace { -/** - * @brief Normalize a graph-route aparam tensor to the flat node axis. - * - * The NeighborGraph ABI carries atomic parameters FLAT on the node axis -- - * shape (N, daparam) with N == n_node_count, the same axis as ``atype`` - * (mirrors ``build_synthetic_graph_inputs`` / ``_build_graph_dynamic_shapes`` - * on the Python export side). The runtime aparam carries the owned (local) - * rows only (``aparam_nall`` is structurally false for pt_expt models, see - * ``init``); on extended-region graphs (multi-rank routes, N == nall_real) - * the ghost rows are zero-padded here. Ghost fitting outputs are never - * retained -- the with-comm artifact masks non-owned energies before - * reduction, and the plain multi-rank remap sums energy over the owned - * prefix only -- so the padded values are inert. - * - * A ghost-only rank (``nlocal == 0``, ``N > 0``) synthesizes a full zero - * tensor: the graph still carries N nodes and the artifact requires the - * (N, daparam) input, while the owned-node mask keeps its contribution - * exactly zero. A missing aparam on a rank that OWNS atoms, or a width - * mismatch, is an explicit error -- the former silently returned the empty - * tensor (artifact reshape failure mid-collective), and a width mismatch - * used to be absorbed by broadcasting ``copy_`` (silent result corruption - * for daparam > 1). - */ -at::Tensor extend_graph_aparam(const at::Tensor& aparam_tensor, - std::int64_t n_node_count, - std::int64_t nlocal, - std::int64_t daparam) { - if (daparam <= 0) { - return aparam_tensor; // model has no aparam input; passed through empty - } - if (aparam_tensor.numel() == 0) { - if (nlocal > 0) { - throw deepmd::deepmd_exception( - "aparam is required (dim_aparam=" + std::to_string(daparam) + - ") but no values were provided on a rank owning " + - std::to_string(nlocal) + " atoms."); - } - // ghost-only rank: there are no owned rows to supply; zeros are inert - // under the owned-node mask but the artifact needs the full node axis. - return torch::zeros({n_node_count, daparam}, aparam_tensor.options()); - } - if (aparam_tensor.numel() != nlocal * daparam) { - throw deepmd::deepmd_exception( - "aparam holds " + std::to_string(aparam_tensor.numel()) + - " values but the graph route expects nlocal * dim_aparam = " + - std::to_string(nlocal) + " * " + std::to_string(daparam) + "."); - } - at::Tensor owned = aparam_tensor.reshape({nlocal, daparam}); - if (nlocal == n_node_count) { - return owned; // single-rank / folded graph: nothing to pad - } - at::Tensor padded = - torch::zeros({n_node_count, daparam}, aparam_tensor.options()); - padded.slice(0, 0, nlocal).copy_(owned); - return padded; -} - -/** - * @brief Assert the flat graph-route aparam contract at the C++ boundary. - * - * The graph artifacts consume aparam FLAT on the node axis, shape - * (N, daparam) -- the layout ``extend_graph_aparam`` produces. A caller - * hand-rolling a rectangular (1, N, daparam) tensor (the pre-flat - * convention) would otherwise fail DEEP inside the artifact -- or, on a - * GPU-only route, only at deployment where no CPU test can catch it (the - * device-edge branch shipped exactly that bug). Failing loudly here turns - * any future such site into an immediate, self-explanatory error. - */ -void check_graph_aparam_flat(const at::Tensor& aparam, - std::int64_t daparam, - const char* where) { - if (daparam <= 0) { - return; - } - if (aparam.dim() != 2 || aparam.size(1) != daparam) { - std::ostringstream oss; - oss << where - << ": graph-route aparam must be flat (N, daparam) on the node axis " - "(produce it with extend_graph_aparam); got a rank-" - << aparam.dim() << " tensor of shape " << aparam.sizes() - << " for daparam = " << daparam << "."; - throw deepmd::deepmd_exception(oss.str()); - } -} +// ``extend_graph_aparam`` and ``check_graph_aparam_flat`` moved to +// commonPT.h (deepmd namespace) so DeepSpinPTExpt.cc can share them for its +// native-spin graph route. void synchronize_current_accelerator_stream() { #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) @@ -325,6 +244,13 @@ void DeepPotPTExpt::init(const std::string& model, lower_input_is_graph_ = false; lower_input_is_canonical_ = false; } + if (lower_input_is_edge_) { + std::cerr << "WARNING: This .pt2 uses the deprecated edge_vec lower " + "schema (pt-backend SeZM/DPA4 freeze). Support will be " + "removed in a future release; refreeze the checkpoint with " + "the pt_expt backend (graph schema)." + << std::endl; + } graph_edge_fp32_ = false; if (metadata.obj_val.count("graph_edge_dtype")) { const std::string graph_edge_dtype = diff --git a/source/api_cc/src/DeepSpin.cc b/source/api_cc/src/DeepSpin.cc index 047b8d85a1..fdaa45af0a 100644 --- a/source/api_cc/src/DeepSpin.cc +++ b/source/api_cc/src/DeepSpin.cc @@ -60,12 +60,13 @@ void DeepSpin::compute(ENERGYTYPE& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam_, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { std::vector dener_; std::vector datom_energy_, datom_virial_; dp->computew(dener_, dforce_, dforce_mag_, dvirial, datom_energy_, datom_virial_, dcoord_, dspin_, datype_, dbox, fparam_, aparam_, - false); + charge_spin, false); dener = dener_[0]; } @@ -79,11 +80,12 @@ void DeepSpin::compute(std::vector& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam_, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { std::vector datom_energy_, datom_virial_; dp->computew(dener, dforce_, dforce_mag_, dvirial, datom_energy_, datom_virial_, dcoord_, dspin_, datype_, dbox, fparam_, aparam_, - false); + charge_spin, false); } // no nlist, no atomic : nframe * precision @@ -96,7 +98,8 @@ template void DeepSpin::compute(ENERGYTYPE& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepSpin::compute(ENERGYTYPE& dener, std::vector& dforce_, @@ -107,7 +110,8 @@ template void DeepSpin::compute(ENERGYTYPE& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepSpin::compute(std::vector& dener, std::vector& dforce_, @@ -118,7 +122,8 @@ template void DeepSpin::compute(std::vector& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepSpin::compute(std::vector& dener, std::vector& dforce_, @@ -129,7 +134,8 @@ template void DeepSpin::compute(std::vector& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); // support spin // nlist, no atomic : nframe @@ -146,12 +152,13 @@ void DeepSpin::compute(ENERGYTYPE& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam_, - const std::vector& aparam__) { + const std::vector& aparam__, + const std::vector& charge_spin) { std::vector dener_; std::vector datom_energy_, datom_virial_; dp->computew(dener_, dforce_, dforce_mag_, dvirial, datom_energy_, datom_virial_, dcoord_, dspin_, datype_, dbox, nghost, lmp_list, - ago, fparam_, aparam__, false); + ago, fparam_, aparam__, charge_spin, false); dener = dener_[0]; } @@ -168,11 +175,12 @@ void DeepSpin::compute(std::vector& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam_, - const std::vector& aparam__) { + const std::vector& aparam__, + const std::vector& charge_spin) { std::vector datom_energy_, datom_virial_; dp->computew(dener, dforce_, dforce_mag_, dvirial, datom_energy_, datom_virial_, dcoord_, dspin_, datype_, dbox, nghost, lmp_list, - ago, fparam_, aparam__, false); + ago, fparam_, aparam__, charge_spin, false); } // nlist, no atomic : nframe * precision @@ -188,7 +196,8 @@ template void DeepSpin::compute(ENERGYTYPE& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_); + const std::vector& aparam_, + const std::vector& charge_spin); template void DeepSpin::compute(ENERGYTYPE& dener, std::vector& dforce_, @@ -202,7 +211,8 @@ template void DeepSpin::compute(ENERGYTYPE& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_); + const std::vector& aparam_, + const std::vector& charge_spin); template void DeepSpin::compute(std::vector& dener, std::vector& dforce_, @@ -216,7 +226,8 @@ template void DeepSpin::compute(std::vector& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_); + const std::vector& aparam_, + const std::vector& charge_spin); template void DeepSpin::compute(std::vector& dener, std::vector& dforce_, @@ -230,7 +241,8 @@ template void DeepSpin::compute(std::vector& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_); + const std::vector& aparam_, + const std::vector& charge_spin); // support spin // no nlist, atomic : nframe @@ -246,11 +258,12 @@ void DeepSpin::compute(ENERGYTYPE& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam_, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { std::vector dener_; dp->computew(dener_, dforce_, dforce_mag_, dvirial, datom_energy_, datom_virial_, dcoord_, dspin_, datype_, dbox, fparam_, aparam_, - true); + charge_spin, true); dener = dener_[0]; } template @@ -265,10 +278,11 @@ void DeepSpin::compute(std::vector& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam_, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { dp->computew(dener, dforce_, dforce_mag_, dvirial, datom_energy_, datom_virial_, dcoord_, dspin_, datype_, dbox, fparam_, aparam_, - true); + charge_spin, true); } // no nlist, atomic : nframe * precision template void DeepSpin::compute(ENERGYTYPE& dener, @@ -282,7 +296,8 @@ template void DeepSpin::compute(ENERGYTYPE& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepSpin::compute(ENERGYTYPE& dener, std::vector& dforce_, @@ -295,7 +310,8 @@ template void DeepSpin::compute(ENERGYTYPE& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepSpin::compute(std::vector& dener, std::vector& dforce_, @@ -308,7 +324,8 @@ template void DeepSpin::compute(std::vector& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepSpin::compute(std::vector& dener, std::vector& dforce_, @@ -321,7 +338,8 @@ template void DeepSpin::compute(std::vector& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); // support spin // nlist, atomic : nframe @@ -340,11 +358,12 @@ void DeepSpin::compute(ENERGYTYPE& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam_, - const std::vector& aparam__) { + const std::vector& aparam__, + const std::vector& charge_spin) { std::vector dener_; dp->computew(dener_, dforce_, dforce_mag_, dvirial, datom_energy_, datom_virial_, dcoord_, dspin_, datype_, dbox, nghost, lmp_list, - ago, fparam_, aparam__, true); + ago, fparam_, aparam__, charge_spin, true); dener = dener_[0]; } template @@ -362,10 +381,11 @@ void DeepSpin::compute(std::vector& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam_, - const std::vector& aparam__) { + const std::vector& aparam__, + const std::vector& charge_spin) { dp->computew(dener, dforce_, dforce_mag_, dvirial, datom_energy_, datom_virial_, dcoord_, dspin_, datype_, dbox, nghost, lmp_list, - ago, fparam_, aparam__, true); + ago, fparam_, aparam__, charge_spin, true); } // nlist, atomic : nframe * precision template void DeepSpin::compute(ENERGYTYPE& dener, @@ -382,7 +402,8 @@ template void DeepSpin::compute(ENERGYTYPE& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_); + const std::vector& aparam_, + const std::vector& charge_spin); template void DeepSpin::compute(ENERGYTYPE& dener, std::vector& dforce_, @@ -398,7 +419,8 @@ template void DeepSpin::compute(ENERGYTYPE& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_); + const std::vector& aparam_, + const std::vector& charge_spin); template void DeepSpin::compute(std::vector& dener, std::vector& dforce_, @@ -414,7 +436,8 @@ template void DeepSpin::compute(std::vector& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_); + const std::vector& aparam_, + const std::vector& charge_spin); template void DeepSpin::compute(std::vector& dener, std::vector& dforce_, @@ -430,7 +453,10 @@ template void DeepSpin::compute(std::vector& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_); + const std::vector& aparam_, + const std::vector& charge_spin); + +int DeepSpin::dim_chg_spin() const { return dp->dim_chg_spin(); } std::vector DeepSpin::get_use_spin() const { if (dp) { @@ -490,7 +516,8 @@ void DeepSpinModelDevi::compute( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { // without nlist if (numb_models == 0) { return; @@ -502,7 +529,7 @@ void DeepSpinModelDevi::compute( for (unsigned ii = 0; ii < numb_models; ++ii) { dps[ii]->compute(all_energy[ii], all_force[ii], all_force_mag[ii], all_virial[ii], dcoord_, dspin_, datype_, dbox, fparam, - aparam_); + aparam_, charge_spin); } } @@ -516,7 +543,8 @@ template void DeepSpinModelDevi::compute( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepSpinModelDevi::compute( std::vector& all_energy, @@ -528,7 +556,8 @@ template void DeepSpinModelDevi::compute( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepSpinModelDevi::compute( @@ -543,7 +572,8 @@ void DeepSpinModelDevi::compute( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { if (numb_models == 0) { return; } @@ -556,7 +586,8 @@ void DeepSpinModelDevi::compute( for (unsigned ii = 0; ii < numb_models; ++ii) { dps[ii]->compute(all_energy[ii], all_force[ii], all_force_mag[ii], all_virial[ii], all_atom_energy[ii], all_atom_virial[ii], - dcoord_, dspin_, datype_, dbox, fparam, aparam_); + dcoord_, dspin_, datype_, dbox, fparam, aparam_, + charge_spin); } } @@ -572,7 +603,8 @@ template void DeepSpinModelDevi::compute( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepSpinModelDevi::compute( std::vector& all_energy, @@ -586,7 +618,8 @@ template void DeepSpinModelDevi::compute( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); // support spin // nlist, no atomic @@ -604,7 +637,8 @@ void DeepSpinModelDevi::compute( const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { if (numb_models == 0) { return; } @@ -615,7 +649,7 @@ void DeepSpinModelDevi::compute( for (unsigned ii = 0; ii < numb_models; ++ii) { dps[ii]->compute(all_energy[ii], all_force[ii], all_force_mag[ii], all_virial[ii], dcoord_, dspin_, datype_, dbox, nghost, - lmp_list, ago, fparam, aparam_); + lmp_list, ago, fparam, aparam_, charge_spin); } } @@ -633,7 +667,8 @@ template void DeepSpinModelDevi::compute( const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepSpinModelDevi::compute( std::vector& all_energy, @@ -648,7 +683,8 @@ template void DeepSpinModelDevi::compute( const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); // support spin // nlist, atomic @@ -668,7 +704,8 @@ void DeepSpinModelDevi::compute( const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { if (numb_models == 0) { return; } @@ -682,7 +719,7 @@ void DeepSpinModelDevi::compute( dps[ii]->compute(all_energy[ii], all_force[ii], all_force_mag[ii], all_virial[ii], all_atom_energy[ii], all_atom_virial[ii], dcoord_, dspin_, datype_, dbox, nghost, lmp_list, ago, - fparam, aparam_); + fparam, aparam_, charge_spin); } } @@ -702,7 +739,8 @@ template void DeepSpinModelDevi::compute( const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepSpinModelDevi::compute( std::vector& all_energy, @@ -719,7 +757,8 @@ template void DeepSpinModelDevi::compute( const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); std::vector DeepSpinModelDevi::get_use_spin() const { if (!dps.empty()) { diff --git a/source/api_cc/src/DeepSpinPTExpt.cc b/source/api_cc/src/DeepSpinPTExpt.cc index c908be18c7..958b240ead 100644 --- a/source/api_cc/src/DeepSpinPTExpt.cc +++ b/source/api_cc/src/DeepSpinPTExpt.cc @@ -2,6 +2,7 @@ #include "DeepSpinPTExpt.h" #if defined(BUILD_PYTORCH) && BUILD_PT_EXPT_SPIN +#include #include #include @@ -102,9 +103,9 @@ void DeepSpinPTExpt::init(const std::string& model, : static_cast(metadata["type_map"].as_array().size()); dfparam = metadata["dim_fparam"].as_int(); daparam = metadata["dim_aparam"].as_int(); - dim_chg_spin = metadata.obj_val.count("dim_chg_spin") - ? metadata["dim_chg_spin"].as_int() - : 0; + dchgspin = metadata.obj_val.count("dim_chg_spin") + ? metadata["dim_chg_spin"].as_int() + : 0; aparam_nall = false; // Spin-specific metadata @@ -143,7 +144,7 @@ void DeepSpinPTExpt::init(const std::string& model, << std::endl; } } - default_chg_spin_ = read_default_chg_spin(metadata, dim_chg_spin); + default_chg_spin_ = read_default_chg_spin(metadata, dchgspin); if (metadata.obj_val.count("do_atomic_virial")) { do_atomic_virial = metadata["do_atomic_virial"].as_bool(); @@ -163,13 +164,35 @@ void DeepSpinPTExpt::init(const std::string& model, } } - // Native spin shares the energy edge ABI; the deepspin scheme keeps the nlist - // contract. Pre-edge spin archives lack the field and default to nlist. + // Native spin shares the energy edge ABI; the deepspin scheme keeps the + // nlist contract; native spin ALSO supports the NeighborGraph schema (no + // dense/nlist lower at all -- see gen_dpa4_spin.py). Pre-edge spin + // archives lack the field and default to nlist. if (metadata.obj_val.count("lower_input_kind")) { - lower_input_is_edge_ = - metadata["lower_input_kind"].as_string() == "edge_vec"; + const std::string lower_input_kind = + metadata["lower_input_kind"].as_string(); + lower_input_is_edge_ = lower_input_kind == "edge_vec"; + lower_input_is_graph_ = lower_input_kind == "graph"; } else { lower_input_is_edge_ = false; + lower_input_is_graph_ = false; + } + if (lower_input_is_edge_) { + std::cerr << "WARNING: This .pt2 uses the deprecated edge_vec lower " + "schema (pt-backend SeZM/DPA4 freeze). Support will be " + "removed in a future release; refreeze the checkpoint with " + "the pt_expt backend (graph schema)." + << std::endl; + } + graph_edge_fp32_ = false; + if (metadata.obj_val.count("graph_edge_dtype")) { + const std::string graph_edge_dtype = + metadata["graph_edge_dtype"].as_string(); + if (graph_edge_dtype != "float32" && graph_edge_dtype != "float64") { + throw deepmd::deepmd_exception( + "metadata graph_edge_dtype must be 'float32' or 'float64'."); + } + graph_edge_fp32_ = graph_edge_dtype == "float32"; } type_map.clear(); @@ -199,6 +222,33 @@ void DeepSpinPTExpt::init(const std::string& model, // pre-PR archives so they retain their previous behaviour. has_message_passing_ = metadata.obj_val.count("has_message_passing") && metadata["has_message_passing"].as_bool(); + + // Model-level pair-type exclusion table -- twin of DeepPotPTExpt::init (see + // there for the full rationale). Exclusion is a BUILD-time transform + // (decision #18/A4): it belongs to the neighbor-graph construction, so the + // C++ ingestion seam applies it exactly once and the exported lower never + // re-applies it. Uploaded once here; UNDEFINED => no exclusion (identity). + { + std::vector> pair_exclude_types; + if (metadata.obj_val.count("pair_exclude_types")) { + for (const auto& v : metadata["pair_exclude_types"].as_array()) { + pair_exclude_types.emplace_back(v[0].as_int(), v[1].as_int()); + } + } + std::vector tbl = + deepmd::buildPairExcludeTable(ntypes, pair_exclude_types); + if (!tbl.empty()) { + torch::Device device(torch::kCUDA, gpu_id); + if (!gpu_enabled) { + device = torch::Device(torch::kCPU); + } + pair_exclude_table_ = + torch::from_blob(tbl.data(), {static_cast(tbl.size())}, + torch::TensorOptions().dtype(torch::kInt32)) + .clone() + .to(device); + } + } if (has_comm_artifact_) { try { with_comm_tempfile_ = std::make_unique( @@ -247,7 +297,8 @@ std::vector DeepSpinPTExpt::run_model( const torch::Tensor& nlist, const torch::Tensor& mapping, const torch::Tensor& fparam, - const torch::Tensor& aparam) { + const torch::Tensor& aparam, + const torch::Tensor& charge_spin) { // Spin model has 7 positional args: coord, atype, spin, nlist, mapping, // fparam, aparam Only include fparam/aparam if the model was exported with // them. @@ -258,11 +309,7 @@ std::vector DeepSpinPTExpt::run_model( if (daparam > 0) { inputs.push_back(aparam); } - if (dim_chg_spin > 0) { - auto charge_spin = torch::tensor(default_chg_spin_, coord.options()) - .view({1, dim_chg_spin}) - .expand({coord.size(0), dim_chg_spin}) - .contiguous(); + if (dchgspin > 0) { inputs.push_back(charge_spin); } return loader->run(inputs); @@ -277,7 +324,8 @@ std::vector DeepSpinPTExpt::run_model_edges( const torch::Tensor& edge_mask, const torch::Tensor& spin, const torch::Tensor& fparam, - const torch::Tensor& aparam) { + const torch::Tensor& aparam, + const torch::Tensor& charge_spin) { // Native-spin edge ABI: the energy edge inputs followed by the // per-local-atom spin leaf, then the optional fparam / aparam / charge_spin. std::vector inputs = { @@ -288,16 +336,131 @@ std::vector DeepSpinPTExpt::run_model_edges( if (daparam > 0) { inputs.push_back(aparam); } - if (dim_chg_spin > 0) { - auto charge_spin = torch::tensor(default_chg_spin_, coord.options()) - .view({1, dim_chg_spin}) - .expand({coord.size(0), dim_chg_spin}) - .contiguous(); + if (dchgspin > 0) { inputs.push_back(charge_spin); } return loader->run(inputs); } +std::vector DeepSpinPTExpt::run_model_graph( + const torch::Tensor& atype, + const torch::Tensor& n_node, + const torch::Tensor& n_local, + const torch::Tensor& edge_index, + const torch::Tensor& edge_vec, + const torch::Tensor& edge_mask, + const torch::Tensor& destination_order, + const torch::Tensor& destination_row_ptr, + const torch::Tensor& source_order, + const torch::Tensor& source_row_ptr, + const torch::Tensor& spin, + const torch::Tensor& fparam, + const torch::Tensor& aparam, + const torch::Tensor& charge_spin) { + // Native-spin graph ABI: the 10 base NeighborGraph tensors, the per-node + // spin leaf (ALWAYS present, positional index 10), then the conditional + // fparam / aparam / charge_spin tail (charge_spin at slot 13 for combined + // native-spin + charge-spin FiLM models, mirroring the energy graph ABI's + // optional charge_spin tail). + deepmd::check_graph_aparam_flat(aparam, daparam, + "DeepSpinPTExpt::run_model_graph"); + std::vector inputs = {atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + spin}; + if (dfparam > 0) { + inputs.push_back(fparam); + } + if (daparam > 0) { + inputs.push_back(aparam); + } + if (dchgspin > 0) { + inputs.push_back(charge_spin); + } + return loader->run(inputs); +} + +std::vector DeepSpinPTExpt::run_model_graph_with_comm( + const torch::Tensor& atype, + const torch::Tensor& n_node, + const torch::Tensor& n_local, + const torch::Tensor& edge_index, + const torch::Tensor& edge_vec, + const torch::Tensor& edge_mask, + const torch::Tensor& destination_order, + const torch::Tensor& destination_row_ptr, + const torch::Tensor& source_order, + const torch::Tensor& source_row_ptr, + const torch::Tensor& spin, + const torch::Tensor& fparam, + const torch::Tensor& aparam, + const torch::Tensor& charge_spin, + const std::vector& comm_tensors) { + if (!with_comm_loader) { + throw deepmd::deepmd_exception( + "run_model_graph_with_comm called but the with-comm artifact is not " + "available. Either the .pt2 has no with-comm artifact compiled " + "(programming error: the caller must check has_comm_artifact_ " + "first), or it failed to load at init (see earlier stderr log). " + "Multi-rank LAMMPS requires a working with-comm artifact."); + } + if (comm_tensors.size() != 8) { + throw deepmd::deepmd_exception( + "run_model_graph_with_comm: comm_tensors must contain exactly 8 " + "tensors (send_list, send_proc, recv_proc, send_num, recv_num, " + "communicator, nlocal, nghost). Got " + + std::to_string(comm_tensors.size()) + "."); + } + deepmd::check_graph_aparam_flat(aparam, daparam, + "DeepSpinPTExpt::run_model_graph_with_comm"); + // Graph-spin with-comm ABI: exactly run_model_graph's prefix (spin stays + // at positional index 10) with the 8 comm tensors appended after the + // conditional fparam/aparam/charge_spin tail -- the twin of + // DeepPotPTExpt::run_model_graph_with_comm, which appends them after the + // energy model's shorter tail. + // + // Device placement follows the energy route: the base tensors (n_local + // included) live on the model device, while ALL 8 comm tensors stay on + // CPU -- border_op's HOST code dereferences their data_ptr and reads + // nlocal/nghost via cheap .item() calls. + // + // ``spin`` is the EXTENDED per-node spin: ghost rows carry their owner's + // spin, delivered by the LAMMPS ``sp`` forward-comm before this call, so + // spin needs no border exchange of its own -- only the per-block ghost + // FEATURE refresh rides border_op. + std::vector inputs = {atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + spin}; + if (dfparam > 0) { + inputs.push_back(fparam); + } + if (daparam > 0) { + inputs.push_back(aparam); + } + if (dchgspin > 0) { + inputs.push_back(charge_spin); + } + for (const auto& ct : comm_tensors) { + inputs.push_back(ct); + } + return with_comm_loader->run(inputs); +} + std::vector DeepSpinPTExpt::run_model_edges_with_comm( const torch::Tensor& coord, const torch::Tensor& atype, @@ -309,6 +472,7 @@ std::vector DeepSpinPTExpt::run_model_edges_with_comm( const torch::Tensor& spin, const torch::Tensor& fparam, const torch::Tensor& aparam, + const torch::Tensor& charge_spin, const std::vector& comm_tensors) { if (!with_comm_loader) { throw deepmd::deepmd_exception( @@ -337,11 +501,7 @@ std::vector DeepSpinPTExpt::run_model_edges_with_comm( if (daparam > 0) { inputs.push_back(aparam); } - if (dim_chg_spin > 0) { - auto charge_spin = torch::tensor(default_chg_spin_, coord.options()) - .view({1, dim_chg_spin}) - .expand({coord.size(0), dim_chg_spin}) - .contiguous(); + if (dchgspin > 0) { inputs.push_back(charge_spin); } for (const auto& t : comm_tensors) { @@ -358,6 +518,7 @@ std::vector DeepSpinPTExpt::run_model_with_comm( const torch::Tensor& mapping, const torch::Tensor& fparam, const torch::Tensor& aparam, + const torch::Tensor& charge_spin, const std::vector& comm_tensors) { if (!with_comm_loader) { throw deepmd::deepmd_exception( @@ -380,11 +541,7 @@ std::vector DeepSpinPTExpt::run_model_with_comm( if (daparam > 0) { inputs.push_back(aparam); } - if (dim_chg_spin > 0) { - auto charge_spin = torch::tensor(default_chg_spin_, coord.options()) - .view({1, dim_chg_spin}) - .expand({coord.size(0), dim_chg_spin}) - .contiguous(); + if (dchgspin > 0) { inputs.push_back(charge_spin); } for (const auto& t : comm_tensors) { @@ -427,6 +584,7 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener, const int& ago, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic) { // Fail fast before allocating any tensors. if (atomic && !do_atomic_virial) { @@ -540,6 +698,21 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener, bool multi_rank = (lmp_list.nprocs > 1); bool atom_map_present = (lmp_list.mapping != nullptr); bool use_with_comm = has_comm_artifact_ && multi_rank; + // NeighborGraph native-spin multi-rank goes through + // ``run_model_graph_with_comm`` (the spin twin of the energy graph + // route). It needs the with-comm artifact for the per-block ghost + // FEATURE refresh; the generic GNN matrix below already fails fast when + // a message-passing model meets multi-rank without one, so the only + // spin-specific guard left is an archive frozen before native spin + // participated in the with-comm export. + if (lower_input_is_graph_ && multi_rank && has_message_passing_ && + !has_comm_artifact_) { + throw deepmd::deepmd_exception( + "multi-rank inference of a graph-kind native-spin .pt2 requires the " + "nested with-comm artifact (has_comm_artifact=false in this " + "archive); re-freeze the model so the with-comm graph lower is " + "compiled, or run on a single MPI rank."); + } // Decision matrix (see PR #5450 description): // non-GNN model (has_message_passing_ == false): regular path is // always safe. @@ -623,6 +796,21 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener, /*fold_to_local=*/!use_with_comm); edge_index_tensor = edge_tensors.edge_index; edge_index_ext_tensor = edge_tensors.edge_index_ext; + } else if (lower_input_is_graph_) { + // Native-spin NeighborGraph route: single-rank folds ghost neighbours + // onto their local owners (``fold_to_local=true``, N == nloc); + // multi-rank indexes the extended node set directly + // (``fold_to_local=false``, N == nall_real) so ghost node features -- + // including the per-node spin embedding -- can be refreshed across + // ranks via border_op (the twin of DeepPotPTExpt.cc's graph branch). + // Cache the skin topology; the model-cutoff edges are recomputed + // on-device every step (see DeepPotPTExpt.cc's graph branch). + const auto edge_tensors = createEdgeTensors( + nlist_data.jlist, dcoord, mapping, nloc, nall_real, device, + /*with_geometry=*/false, /*row_centers=*/&nlist_data.ilist, + /*fold_to_local=*/!use_with_comm); + edge_index_tensor = edge_tensors.edge_index; + edge_index_ext_tensor = edge_tensors.edge_index_ext; } else { nlist_data.padding(); // Flatten raw nlist — the .pt2 model sorts by distance on-device. @@ -680,6 +868,40 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener, aparam_tensor = torch::zeros({0}, options).to(device); } + // Build charge_spin tensor: use the runtime value when provided, fall back + // to default_chg_spin_ stored in the .pt2 metadata. Mirrors + // DeepPotPTExpt::compute -- these spin paths are single-frame, so the + // runtime vector must hold exactly dim_chg_spin values. + at::Tensor charge_spin_tensor; + if (dchgspin > 0) { + auto dbl_options = torch::TensorOptions().dtype(torch::kFloat64); + if (!charge_spin.empty()) { + if (static_cast(charge_spin.size()) != dchgspin) { + throw deepmd::deepmd_exception( + "charge_spin has " + std::to_string(charge_spin.size()) + + " values but the model expects dim_chg_spin=" + + std::to_string(dchgspin) + "."); + } + charge_spin_tensor = + torch::from_blob(const_cast(charge_spin.data()), + {1, static_cast(charge_spin.size())}, + dbl_options) + .clone() + .to(device); + } else if (!default_chg_spin_.empty()) { + charge_spin_tensor = + torch::from_blob(const_cast(default_chg_spin_.data()), + {1, dchgspin}, dbl_options) + .clone() + .to(device); + } else { + throw deepmd::deepmd_exception( + "charge_spin is empty and no default_chg_spin is available in the " + ".pt2 metadata. Provide charge_spin explicitly or regenerate the " + "model with a default charge/spin value."); + } + } + // Phase 4 dispatch: route to with-comm artifact in multi-rank mode. // ``has_spin=tensor([1])`` is baked into the with-comm graph at // trace time (Phase 3, spin_model.forward_common_lower_exportable @@ -744,7 +966,87 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener, lmp_list, lmp_list.sendlist, lmp_list.sendnum, lmp_list.recvnum, nloc, nghost_real); } - if (lower_input_is_edge_) { + if (lower_input_is_graph_) { + // Native-spin NeighborGraph multi-rank: the twin of + // DeepPotPTExpt's graph with-comm branch, plus the EXTENDED per-node + // spin. Nodes are the extended set (fold_to_local=false above), so + // ``n_node`` counts owned+ghost while the separate device + // ``n_local`` drives the owned-energy mask; ghost spins arrive via + // the LAMMPS ``sp`` forward-comm, and border_op refreshes ghost node + // FEATURES between interaction blocks. + // Collective preflight (twin of DeepPotPTExpt.cc's graph with-comm + // branch): a rank with zero owned+ghost atoms cannot run the graph + // artifact, and a rank-LOCAL throw would leave the non-empty peers + // blocked forever in the per-layer border_op collectives. All-reduce + // the minimum node count over the LAMMPS communicator + // (``comm_tensors[5]``) so EVERY rank agrees to run -- or every rank + // throws promptly with the same error. The op is an identity when the + // communicator handle is null or MPI is not compiled in. + // + // Cached across ``ago > 0`` force calls: the owned+ghost node count + // shares the lifetime of the cached nlist/mapping/edge topology (both + // only change on an ``ago == 0`` rebuild, which is globally + // synchronized by LAMMPS), so re-running the collective on every + // cache-hit MD step would add a global synchronization to the hot path + // with no added protection. + if (ago == 0 || !graph_comm_preflight_done_) { + graph_comm_preflight_done_ = false; + const auto allreduce_min = + c10::Dispatcher::singleton() + .findSchemaOrThrow("deepmd_export::allreduce_min_int", "") + .typed(); + at::Tensor local_n_node = + torch::full({1}, static_cast(nall_real), int_option); + const std::int64_t global_min_n_node = + allreduce_min.call(local_n_node, comm_tensors[5].to(torch::kCPU)) + .item(); + if (global_min_n_node <= 0) { + throw deepmd::deepmd_exception( + "Multi-rank native-spin graph inference does not support a rank " + "with zero owned+ghost atoms (the exported graph artifact needs " + "at least one node, and skipping the run would desync the " + "per-layer MPI ghost exchange; this rank has " + + std::to_string(nall_real) + + " owned+ghost atoms). Use a domain decomposition that keeps " + "every rank non-empty, or a dense .pt2."); + } + graph_comm_preflight_done_ = true; + } + const auto edge_tensors = + compactEdgeTensors(edge_index_tensor, edge_index_ext_tensor, + coord_Tensor, static_cast(rcut)); + const std::int64_t n_node_count = nall_real; + at::Tensor n_node_tensor = + torch::full({1}, n_node_count, int_option).to(device); + at::Tensor n_local_tensor = + torch::full({1}, static_cast(nloc), int_option) + .to(device); + at::Tensor node_atype = + atype_Tensor.slice(1, 0, n_node_count).reshape({n_node_count}); + GraphTensorPack graph_pack; + graph_pack.atype = node_atype; + graph_pack.n_node = n_node_tensor; + graph_pack.n_local = n_local_tensor; + graph_pack.edge_index = edge_tensors.edge_index; + graph_pack.edge_vec = graph_edge_fp32_ + ? edge_tensors.edge_vec.to(torch::kFloat32) + : edge_tensors.edge_vec; + // Same build-time exclusion seam as the single-rank graph branch. + graph_pack.edge_mask = deepmd::applyPairExclusion( + edge_tensors.edge_index, edge_tensors.edge_mask, node_atype, + pair_exclude_table_, ntypes); + canonicalizeGraphPayload(graph_pack, n_node_count); + flat_outputs = run_model_graph_with_comm( + node_atype, n_node_tensor, n_local_tensor, graph_pack.edge_index, + graph_pack.edge_vec, graph_pack.edge_mask, + graph_pack.destination_order, graph_pack.destination_row_ptr, + graph_pack.source_order, graph_pack.source_row_ptr, + spin_Tensor.slice(1, 0, n_node_count).reshape({n_node_count, 3}), + fparam_tensor, + deepmd::extend_graph_aparam(aparam_tensor, n_node_count, nloc, + daparam), + charge_spin_tensor, comm_tensors); + } else if (lower_input_is_edge_) { // Native spin multi-rank: edges index the extended node set // (fold_to_local=false above), the EXTENDED per-node spin feeds the // descriptor (ghost spins arrive via the LAMMPS sp forward-comm), and @@ -764,7 +1066,8 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener, flat_outputs = run_model_edges_with_comm( coord_Tensor, atype_Tensor.slice(1, 0, nloc), atype_Tensor, ph_edge_index, ph_edge_vec, ph_edge_index, ph_edge_mask, - spin_Tensor, fparam_tensor, aparam_tensor, comm_tensors); + spin_Tensor, fparam_tensor, aparam_tensor, charge_spin_tensor, + comm_tensors); } else { const auto edge_tensors = compactEdgeTensors(edge_index_tensor, edge_index_ext_tensor, @@ -773,13 +1076,88 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener, coord_Tensor, atype_Tensor.slice(1, 0, nloc), atype_Tensor, edge_tensors.edge_index, edge_tensors.edge_vec, edge_tensors.edge_index_ext, edge_tensors.edge_mask, spin_Tensor, - fparam_tensor, aparam_tensor, comm_tensors); + fparam_tensor, aparam_tensor, charge_spin_tensor, comm_tensors); } } else { + // Model-level pair exclusion is a BUILD-time transform (decision + // #18/A4): the exported dense lower consumes a pre-excluded nlist and + // never re-applies it. The multi-rank (with-comm) dense route shares + // the same dense nlist as the single-rank path below, so it applies the + // SAME seam -- otherwise a message-passing spin .pt2 with + // pair_exclude_types would silently include excluded pairs on the + // with-comm path. The cross-rank ghost exchange happens inside + // run_model_with_comm and does not change the nlist's meaning, so + // pre-excluding it is correct per rank. ``atype_Tensor`` is the + // real-atom (and, on an empty subdomain, phantom-prefixed) extended + // type vector that ``firstneigh_tensor`` indexes, so both live in the + // same index space; the spin model's internal atom doubling happens + // downstream of this seam. + const at::Tensor excl_nlist = deepmd::applyPairExclusionNlist( + firstneigh_tensor, atype_Tensor, pair_exclude_table_, ntypes); flat_outputs = run_model_with_comm( - coord_Tensor, atype_Tensor, spin_Tensor, firstneigh_tensor, - mapping_tensor, fparam_tensor, aparam_tensor, comm_tensors); + coord_Tensor, atype_Tensor, spin_Tensor, excl_nlist, mapping_tensor, + fparam_tensor, aparam_tensor, charge_spin_tensor, comm_tensors); + } + } else if (lower_input_is_graph_) { + if (nall_real == 0) { + // Truly-empty rank (no real local atoms AND no real ghosts): the graph + // would emit N == 0 nodes, which violates the exported + // ``Dim("n_node_total", min=1)``. Such a rank contributes nothing, so + // fill zero outputs and return instead of running the artifact. Twin of + // DeepPotPTExpt::compute_inner's non-comm graph guard. (The + // ``nloc_real == 0`` empty-subdomain case has ``nall_real > 0`` -- real + // ghosts within rcut -- so it is phantom-padded above and still runs the + // model normally.) + ener.assign(nframes, static_cast(0)); + force.assign(static_cast(nframes) * fwd_map.size() * 3, + static_cast(0)); + force_mag.assign(static_cast(nframes) * fwd_map.size() * 3, + static_cast(0)); + virial.assign(static_cast(nframes) * 9, + static_cast(0)); + if (atomic) { + atom_energy.assign(static_cast(nframes) * fwd_map.size(), + static_cast(0)); + atom_virial.assign(static_cast(nframes) * fwd_map.size() * 9, + static_cast(0)); + } + return; } + // Native-spin NeighborGraph route: single-rank ONLY (guaranteed by the + // multi-rank fail-fast above). Compact the cached skin topology to the + // model cutoff and feed the OWNED-atom spin (nloc, 3) -- ghosts are + // already folded onto their local owners via edge_index + // (fold_to_local=true above), so no separate ghost spin node is needed. + const auto edge_tensors = + compactEdgeTensors(edge_index_tensor, edge_index_ext_tensor, + coord_Tensor, static_cast(rcut)); + at::Tensor n_node_tensor = + torch::full({1}, static_cast(nloc), int_option) + .to(device); + at::Tensor node_atype = atype_Tensor.slice(1, 0, nloc).reshape({nloc}); + GraphTensorPack graph_pack; + graph_pack.atype = node_atype; + graph_pack.n_node = n_node_tensor; + graph_pack.n_local = n_node_tensor; + graph_pack.edge_index = edge_tensors.edge_index; + graph_pack.edge_vec = graph_edge_fp32_ + ? edge_tensors.edge_vec.to(torch::kFloat32) + : edge_tensors.edge_vec; + // Model-level pair exclusion belongs to the graph BUILD (decision #18/A4), + // exactly as on the non-spin route: the exported lower consumes a + // pre-excluded edge_mask and never re-applies it. + graph_pack.edge_mask = deepmd::applyPairExclusion( + edge_tensors.edge_index, edge_tensors.edge_mask, node_atype, + pair_exclude_table_, ntypes); + canonicalizeGraphPayload(graph_pack, nloc); + flat_outputs = run_model_graph( + node_atype, n_node_tensor, n_node_tensor, graph_pack.edge_index, + graph_pack.edge_vec, graph_pack.edge_mask, graph_pack.destination_order, + graph_pack.destination_row_ptr, graph_pack.source_order, + graph_pack.source_row_ptr, + spin_Tensor.slice(1, 0, nloc).reshape({nloc, 3}), fparam_tensor, + deepmd::extend_graph_aparam(aparam_tensor, nloc, nloc, daparam), + charge_spin_tensor); } else if (lower_input_is_edge_) { // Native spin edge path (single-rank): recompute the model-cutoff edge // vectors from the cached skin topology and feed only the owned-atom @@ -793,16 +1171,42 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener, coord_Tensor, atype_Tensor.slice(1, 0, nloc), edge_tensors.edge_index, edge_tensors.edge_vec, edge_tensors.edge_index_ext, edge_tensors.edge_mask, spin_Tensor.slice(1, 0, nloc), fparam_tensor, - aparam_tensor); + aparam_tensor, charge_spin_tensor); } else { - flat_outputs = - run_model(coord_Tensor, atype_Tensor, spin_Tensor, firstneigh_tensor, - mapping_tensor, fparam_tensor, aparam_tensor); + // Model-level pair exclusion is a BUILD-time transform (decision #18/A4): + // the exported dense lower consumes a pre-excluded nlist and never + // re-applies it. Single-rank dense application site; the multi-rank + // (with-comm) dense sibling above applies the same seam. + const at::Tensor excl_nlist = deepmd::applyPairExclusionNlist( + firstneigh_tensor, atype_Tensor, pair_exclude_table_, ntypes); + flat_outputs = run_model(coord_Tensor, atype_Tensor, spin_Tensor, + excl_nlist, mapping_tensor, fparam_tensor, + aparam_tensor, charge_spin_tensor); } std::map output_map; extract_outputs(output_map, flat_outputs); + if (lower_input_is_graph_) { + // The graph forward emits flat-N PUBLIC keys (atom_energy/energy/force/ + // force_mag/virial/atom_virial); rewrite them into the dense internal-key + // layout the shared extraction below expects. + // + // The two node layouts need DIFFERENT remaps, and picking the wrong one + // is not a silent error: single-rank folds ghosts onto owners + // (fold_to_local=true, N == nloc) so the per-node outputs are padded up + // to nall, while the with-comm route keeps the extended node set + // (fold_to_local=false, N == nall) and must NOT pad -- padding there + // throws on the index_put_ as soon as nloc < nall. + if (use_with_comm) { + deepmd::remap_graph_spin_outputs_to_dense_keys_extended( + output_map, nloc, nall_real, atomic); + } else { + deepmd::remap_graph_spin_outputs_to_dense_keys(output_map, nloc, + nall_real, atomic); + } + } + // Extract energy torch::Tensor flat_energy_ = output_map["energy_redu"].view({-1}).to(torch::kCPU); @@ -918,6 +1322,7 @@ template void DeepSpinPTExpt::compute>( const int& ago, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); template void DeepSpinPTExpt::compute>( std::vector& ener, @@ -935,6 +1340,7 @@ template void DeepSpinPTExpt::compute>( const int& ago, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); // ============================================================================ @@ -954,6 +1360,7 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener, const std::vector& box, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic) { // Fail fast before allocating any tensors. if (atomic && !do_atomic_virial) { @@ -1076,11 +1483,20 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener, .to(device); at::Tensor nlist_tensor; EdgeTensorPack edge_tensors; + GraphTensorPack graph_tensors; if (lower_input_is_edge_) { // Native spin edge ABI: build the full edge schema once (no cached skin // topology in the standalone path), folding ghosts onto local owners. edge_tensors = createEdgeTensors(nlist_raw, coord_cpy_d, mapping_64, nloc, nall, device); + } else if (lower_input_is_graph_) { + // Standalone (no nlist) graph schema: build_nlist already cut at rcut + // and keys row i to center i, so no row_centers remapping is needed. + // Single-rank only (the standalone build_nlist path never sees a + // multi-rank comm), so fold_to_local defaults to true (N == nloc). + graph_tensors = + buildGraphTensors(nlist_raw, coord_cpy_d, atype_cpy, mapping_64, nloc, + nall, static_cast(rcut), device); } else { // Flatten raw nlist — the .pt2 model sorts by distance on-device. nlist_tensor = @@ -1127,25 +1543,93 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener, aparam_tensor = torch::zeros({0}, options).to(device); } + // Build charge_spin tensor: use the runtime value when provided, fall back + // to default_chg_spin_ stored in the .pt2 metadata. Mirrors + // DeepPotPTExpt::compute -- these spin paths are single-frame, so the + // runtime vector must hold exactly dim_chg_spin values. + at::Tensor charge_spin_tensor; + if (dchgspin > 0) { + auto dbl_options = torch::TensorOptions().dtype(torch::kFloat64); + if (!charge_spin.empty()) { + if (static_cast(charge_spin.size()) != dchgspin) { + throw deepmd::deepmd_exception( + "charge_spin has " + std::to_string(charge_spin.size()) + + " values but the model expects dim_chg_spin=" + + std::to_string(dchgspin) + "."); + } + charge_spin_tensor = + torch::from_blob(const_cast(charge_spin.data()), + {1, static_cast(charge_spin.size())}, + dbl_options) + .clone() + .to(device); + } else if (!default_chg_spin_.empty()) { + charge_spin_tensor = + torch::from_blob(const_cast(default_chg_spin_.data()), + {1, dchgspin}, dbl_options) + .clone() + .to(device); + } else { + throw deepmd::deepmd_exception( + "charge_spin is empty and no default_chg_spin is available in the " + ".pt2 metadata. Provide charge_spin explicitly or regenerate the " + "model with a default charge/spin value."); + } + } + // 5. Run the .pt2 model: native spin uses the energy edge ABI plus the - // owned-atom spins; the deepspin scheme keeps the 7-arg nlist contract. + // owned-atom spins; the deepspin scheme keeps the 7-arg nlist contract; + // the NeighborGraph route runs the graph artifact with the owned-atom + // spin (nloc, 3) as its 11th positional input. std::vector flat_outputs; if (lower_input_is_edge_) { flat_outputs = run_model_edges( coord_Tensor, atype_Tensor.slice(1, 0, nloc), edge_tensors.edge_index, edge_tensors.edge_vec, edge_tensors.edge_index_ext, edge_tensors.edge_mask, spin_Tensor.slice(1, 0, nloc), fparam_tensor, - aparam_tensor); + aparam_tensor, charge_spin_tensor); + } else if (lower_input_is_graph_) { + // Same build-time seam as the cached-nlist branch above. + graph_tensors.edge_mask = deepmd::applyPairExclusion( + graph_tensors.edge_index, graph_tensors.edge_mask, graph_tensors.atype, + pair_exclude_table_, ntypes); + canonicalizeGraphPayload(graph_tensors, graph_tensors.atype.size(0)); + if (graph_edge_fp32_) { + graph_tensors.edge_vec = graph_tensors.edge_vec.to(torch::kFloat32); + } + flat_outputs = run_model_graph( + graph_tensors.atype, graph_tensors.n_node, graph_tensors.n_local, + graph_tensors.edge_index, graph_tensors.edge_vec, + graph_tensors.edge_mask, graph_tensors.destination_order, + graph_tensors.destination_row_ptr, graph_tensors.source_order, + graph_tensors.source_row_ptr, + spin_Tensor.slice(1, 0, nloc).reshape({nloc, 3}), fparam_tensor, + deepmd::extend_graph_aparam(aparam_tensor, natoms, natoms, daparam), + charge_spin_tensor); } else { - flat_outputs = - run_model(coord_Tensor, atype_Tensor, spin_Tensor, nlist_tensor, - mapping_tensor, fparam_tensor, aparam_tensor); + // Model-level pair exclusion is a BUILD-time transform (decision #18/A4): + // the exported dense lower consumes a pre-excluded nlist and never + // re-applies it; this is the single application site on the standalone + // (build_nlist) dense route. + const at::Tensor excl_nlist = deepmd::applyPairExclusionNlist( + nlist_tensor, atype_Tensor, pair_exclude_table_, ntypes); + flat_outputs = run_model(coord_Tensor, atype_Tensor, spin_Tensor, + excl_nlist, mapping_tensor, fparam_tensor, + aparam_tensor, charge_spin_tensor); } // 6. Extract outputs std::map output_map; extract_outputs(output_map, flat_outputs); + if (lower_input_is_graph_) { + // The graph forward emits LOCAL public keys; rewrite them into the dense + // internal-key layout used below. nloc == N (graph node count); the + // standalone (build_nlist) path is always single-rank. + deepmd::remap_graph_spin_outputs_to_dense_keys(output_map, nloc, nall, + atomic); + } + // 7. Extract energy torch::Tensor flat_energy_ = output_map["energy_redu"].view({-1}).to(torch::kCPU); @@ -1212,6 +1696,7 @@ template void DeepSpinPTExpt::compute>( const std::vector& box, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); template void DeepSpinPTExpt::compute>( std::vector& ener, @@ -1226,6 +1711,7 @@ template void DeepSpinPTExpt::compute>( const std::vector& box, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); void DeepSpinPTExpt::get_type_map(std::string& type_map_str) { @@ -1254,7 +1740,46 @@ void DeepSpinPTExpt::computew(std::vector& ener, const bool atomic) { translate_error([&] { compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord, - spin, atype, box, fparam, aparam, atomic); + spin, atype, box, fparam, aparam, {}, atomic); + }); +} +void DeepSpinPTExpt::computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic) { + translate_error([&] { + compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord, + spin, atype, box, fparam, aparam, {}, atomic); + }); +} +void DeepSpinPTExpt::computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic) { + translate_error([&] { + compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord, + spin, atype, box, nghost, inlist, ago, fparam, aparam, {}, atomic); }); } void DeepSpinPTExpt::computew(std::vector& ener, @@ -1267,12 +1792,36 @@ void DeepSpinPTExpt::computew(std::vector& ener, const std::vector& spin, const std::vector& atype, const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, const std::vector& fparam, const std::vector& aparam, const bool atomic) { translate_error([&] { compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord, - spin, atype, box, fparam, aparam, atomic); + spin, atype, box, nghost, inlist, ago, fparam, aparam, {}, atomic); + }); +} + +// forward to template method (runtime charge_spin) +void DeepSpinPTExpt::computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + translate_error([&] { + compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord, + spin, atype, box, fparam, aparam, charge_spin, atomic); }); } void DeepSpinPTExpt::computew(std::vector& ener, @@ -1290,10 +1839,33 @@ void DeepSpinPTExpt::computew(std::vector& ener, const int& ago, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + translate_error([&] { + compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord, + spin, atype, box, nghost, inlist, ago, fparam, aparam, charge_spin, + atomic); + }); +} + +// forward to template method (runtime charge_spin) +void DeepSpinPTExpt::computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, const bool atomic) { translate_error([&] { compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord, - spin, atype, box, nghost, inlist, ago, fparam, aparam, atomic); + spin, atype, box, fparam, aparam, charge_spin, atomic); }); } void DeepSpinPTExpt::computew(std::vector& ener, @@ -1311,10 +1883,13 @@ void DeepSpinPTExpt::computew(std::vector& ener, const int& ago, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic) { translate_error([&] { compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord, - spin, atype, box, nghost, inlist, ago, fparam, aparam, atomic); + spin, atype, box, nghost, inlist, ago, fparam, aparam, charge_spin, + atomic); }); } + #endif diff --git a/source/api_cc/tests/test_deeppot_dpa4_zbl_ptexpt.cc b/source/api_cc/tests/test_deeppot_dpa4_zbl_ptexpt.cc new file mode 100644 index 0000000000..8dfe1efaaf --- /dev/null +++ b/source/api_cc/tests/test_deeppot_dpa4_zbl_ptexpt.cc @@ -0,0 +1,220 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// C++ inference for a DPA4 model with analytical ZBL bridging (pt_expt +// .pt2, graph lower). +// +// ``bridging_method: ZBL`` builds a COMPOSITION -- LinearEnergyModel over +// [learned DPA4, InterPotentialAtomicModel] -- so this exercises the graph +// lower of a linear composition, which no other C++ fixture covers. Before +// this test ZBL bridging had NO C++ or LAMMPS coverage at all: its only +// end-to-end check drove the archive through the PYTHON DeepPot, which +// never reaches DeepPotPTExpt. +// +// Single-rank only by construction: bridging enables the descriptor's +// Source Freeze Propagation Gate, whose per-node eta folds a node's full +// outgoing-edge set, and edges exist only for owned centres -- so the +// generator asserts has_comm_artifact=false and multi-rank is out of scope +// here. +// +// The references come from the Python DeepEval of the SAME archive +// (source/tests/infer/gen_dpa4_zbl.py), so a 1e-10 match validates the +// whole C++ chain: metadata parse, graph ingestion, and the compiled math +// of the composition. +#include + +#include +#include +#include +#include + +#include "DeepPot.h" +// Defines BUILD_PT_EXPT (via its __has_include probe for the inductor +// headers). Without this include the guards below see it undefined and +// every case GTEST_SKIPs with "PyTorch support is not enabled" -- silently, +// which is how this suite first ran as 8 skips that looked like passes. +#include "DeepPotPTExpt.h" +#include "expected_ref.h" +#include "neighbor_list.h" +#include "test_utils.h" + +namespace { +constexpr const char* kModelPath = + "../../tests/infer/deeppot_dpa4_zbl_graph.pt2"; +constexpr const char* kRefPath = + "../../tests/infer/deeppot_dpa4_zbl_graph.expected"; + +// Magnitude-scaled bound. The analytical ZBL term on a 0.9 A pair makes +// this fixture's forces ~1.4e3 -- deliberately, so the term cannot be +// mistaken for noise -- and ONE fp32 ULP at that magnitude is 2^-13 = +// 1.22e-4. The suite's flat 1e-4 float bound therefore asks for sub-ULP +// agreement, which no fp32 result can satisfy: the observed deltas were +// exact binary fractions (2^-13, 3*2^-14, 6*2^-14), i.e. 1-3 ULP of the +// representation, not numerical error. fp64 keeps the strict absolute +// 1e-10 (one fp64 ULP here is 2.3e-13, so 1e-10 is still ~400x tighter +// than the representation). +template +inline double zbl_tol(double expected) { + if (std::is_same::value) { + return 1e-10; + } + // 5e-7 relative is ~4 fp32 ULP; the 1e-4 floor keeps near-zero + // components at the suite's usual float bound. + return 1e-4 + 5e-7 * fabs(expected); +} +} // namespace + +template +class TestInferDeepPotDpa4ZblPtExpt : public ::testing::Test { + protected: + // Fixed 6-atom system -- verbatim from gen_dpa4_zbl.py's _COORDS/_ATYPES. + // Atoms 0 and 1 sit 0.9 A apart, inside bridging_r_outer, so the + // analytical ZBL term dominates their interaction. + std::vector coord = {1.0, 1.0, 1.0, 1.9, 1.0, 1.0, 1.3, 1.8, 1.0, + 0.4, 1.2, 1.6, 3.6, 2.0, 1.3, 3.4, 0.7, 1.7}; + std::vector atype = {0, 0, 0, 1, 1, 1}; + std::vector box = {6., 0., 0., 0., 6., 0., 0., 0., 6.}; + + std::vector expected_e; + std::vector expected_f; + std::vector expected_tot_v; + std::vector expected_e_nopbc; + std::vector expected_f_nopbc; + + int natoms; + double expected_tot_e; + double expected_tot_e_nopbc; + + deepmd::DeepPot dp; + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + std::ifstream model_file(kModelPath); + if (!model_file.good()) { + GTEST_SKIP() << "Skip because " << kModelPath + << " was not generated (run " + "source/tests/infer/gen_dpa4_zbl.py)."; + } + dp.init(kModelPath); + + deepmd_test::ExpectedRef ref; + ref.load(kRefPath); + expected_e = ref.get("pbc", "expected_e"); + expected_f = ref.get("pbc", "expected_f"); + expected_tot_v = ref.get("pbc", "expected_tot_v"); + expected_e_nopbc = ref.get("nopbc", "expected_e"); + expected_f_nopbc = ref.get("nopbc", "expected_f"); + + natoms = expected_e.size(); + EXPECT_EQ(natoms * 3, expected_f.size()); + EXPECT_EQ(9, expected_tot_v.size()); + expected_tot_e = 0.; + expected_tot_e_nopbc = 0.; + for (int ii = 0; ii < natoms; ++ii) { + expected_tot_e += expected_e[ii]; + expected_tot_e_nopbc += expected_e_nopbc[ii]; + } + }; + + void TearDown() override {}; +}; + +TYPED_TEST_SUITE(TestInferDeepPotDpa4ZblPtExpt, ValueTypes); + +TYPED_TEST(TestInferDeepPotDpa4ZblPtExpt, type_map) { + std::string type_map; + this->dp.get_type_map(type_map); + EXPECT_EQ(type_map, "Ni O"); +} + +// Standalone path (no InputNlist): DeepPotPTExpt::compute -> the graph +// branch, on a linear composition. +TYPED_TEST(TestInferDeepPotDpa4ZblPtExpt, cpu_build_nlist) { + using VALUETYPE = TypeParam; + const std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_f = this->expected_f; + std::vector& expected_tot_v = this->expected_tot_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + + double ener; + std::vector force, virial; + this->dp.compute(ener, force, virial, coord, atype, box); + + EXPECT_EQ(force.size(), natoms * 3); + // anti-vacuity: the close Ni-Ni pair must drive a large ZBL repulsion, + // else the fixture would be degenerate and this comparison meaningless. + double fmax = 0.; + for (int ii = 0; ii < natoms * 3; ++ii) { + fmax = std::max(fmax, static_cast(fabs(force[ii]))); + } + EXPECT_GT(fmax, 1e-3) << "forces are trivially small; fixture is vacuous"; + + EXPECT_LT(fabs(ener - expected_tot_e), zbl_tol(expected_tot_e)); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), + zbl_tol(expected_f[ii])); + } + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), + zbl_tol(expected_tot_v[ii])); + } +} + +TYPED_TEST(TestInferDeepPotDpa4ZblPtExpt, cpu_build_nlist_atomic) { + using VALUETYPE = TypeParam; + const std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + + double ener; + std::vector force, virial, atom_ener, atom_vir; + this->dp.compute(ener, force, virial, atom_ener, atom_vir, coord, atype, box); + + EXPECT_EQ(atom_ener.size(), natoms); + EXPECT_LT(fabs(ener - expected_tot_e), zbl_tol(expected_tot_e)); + for (int ii = 0; ii < natoms; ++ii) { + EXPECT_LT(fabs(atom_ener[ii] - expected_e[ii]), + zbl_tol(expected_e[ii])); + } + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), + zbl_tol(expected_f[ii])); + } +} + +// LAMMPS path (explicit InputNlist, nghost=0): the gas-phase system, so it +// is compared against the NoPBC reference. +TYPED_TEST(TestInferDeepPotDpa4ZblPtExpt, cpu_lmp_nlist) { + using VALUETYPE = TypeParam; + const std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& expected_f = this->expected_f_nopbc; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e_nopbc; + std::vector box = {}; + + double ener; + std::vector force, virial; + std::vector > nlist_data = { + {1, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, {0, 1, 3, 4, 5}, + {0, 1, 2, 4, 5}, {0, 1, 2, 3, 5}, {0, 1, 2, 3, 4}}; + std::vector ilist(natoms), numneigh(natoms); + std::vector firstneigh(natoms); + deepmd::InputNlist inlist(natoms, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + this->dp.compute(ener, force, virial, coord, atype, box, 0, inlist, 0); + + EXPECT_LT(fabs(ener - expected_tot_e), zbl_tol(expected_tot_e)); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), + zbl_tol(expected_f[ii])); + } +} diff --git a/source/api_cc/tests/test_deeppot_universal.cc b/source/api_cc/tests/test_deeppot_universal.cc index b4482868f1..964ca9fd4f 100644 --- a/source/api_cc/tests/test_deeppot_universal.cc +++ b/source/api_cc/tests/test_deeppot_universal.cc @@ -325,7 +325,30 @@ std::vector variant_deeppot_cases() { /*supports_no_pbc_simple=*/true, /*supports_no_pbc_atomic=*/false, /*supports_no_pbc_lmp_nlist=*/true, - /*supports_no_pbc_lmp_nlist_atomic=*/false}}; + /*supports_no_pbc_lmp_nlist_atomic=*/false}, + {"dpa4_graph_pytorch_pt2", + Backend::PTExpt, + "../../tests/infer/deeppot_dpa4_graph.pt2", + /*convert_pbtxt=*/false, + nullptr, + nullptr, + "../../tests/infer/deeppot_dpa4_graph.expected", + "pbc", + "nopbc", + 1e-10, + 1e-4, + /*supports_float=*/true, + /*supports_finite_difference=*/true, + /*supports_lmp_nlist=*/true, + /*supports_lmp_nlist_atomic=*/true, + /*supports_lmp_nlist_cutoff_twice=*/true, + /*supports_lmp_nlist_type_sel=*/true, + /*supports_print_summary=*/true, + /*supports_no_pbc_simple=*/true, + /*supports_no_pbc_atomic=*/false, + /*supports_no_pbc_lmp_nlist=*/true, + /*supports_no_pbc_lmp_nlist_atomic=*/false, + /*skip_if_artifact_missing=*/true}}; } std::vector default_fparam_cases() { @@ -1717,8 +1740,14 @@ TEST_P(VariantDeepPotTest, FiniteDifferenceFloat) { GTEST_SKIP() << GetParam().name << " finite-difference coverage is not enabled."; } + // DPA4 computes in reduced precision (same established bound for both the + // dense-nlist and graph lowers -- they share the same descriptor/fitting + // math, only the lower schema differs). const double finite_difference_tol = - GetParam().name == "dpa4_pytorch_pt2" ? 3e-2 : -1.0; + (GetParam().name == "dpa4_pytorch_pt2" || + GetParam().name == "dpa4_graph_pytorch_pt2") + ? 3e-2 + : -1.0; check_finite_difference(dp, finite_difference_tol); } diff --git a/source/api_cc/tests/test_deepspin_dpa4_chgspin_ptexpt.cc b/source/api_cc/tests/test_deepspin_dpa4_chgspin_ptexpt.cc new file mode 100644 index 0000000000..d9cfd743ba --- /dev/null +++ b/source/api_cc/tests/test_deepspin_dpa4_chgspin_ptexpt.cc @@ -0,0 +1,373 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Functional test for the RUNTIME charge_spin argument of the DeepSpin +// (.pt2 / pt_expt) inference path. +// +// Every other spin fixture in this tree has dim_chg_spin == 0, so the +// charge_spin argument threaded through DeepSpin::compute -> DeepSpinPTExpt +// is inert on them: a dead ingestion seam would pass all of them. The model +// here (source/tests/infer/gen_dpa4_spin_chgspin.py) is the only one that is +// BOTH is_spin and dim_chg_spin == 2 -- a native-spin DPA4 whose descriptor +// also carries add_chg_spin_ebd=True -- so the argument is load-bearing. +// +// Asserted here: +// * two DIFFERENT charge_spin vectors give DIFFERENT energies, each +// matching its own reference section (the seam is live AND correct); +// * an EMPTY charge_spin reproduces the model's stored default_chg_spin +// (backward compatibility -- the property most likely to regress, since +// every pre-existing caller passes nothing); +// * passing the default value explicitly equals passing nothing (the two +// ways of selecting the default agree); +// * both DeepSpinPTExpt::compute overloads are covered: the standalone +// (build-nlist) one and the LAMMPS (InputNlist) one, each of which does +// its own charge_spin -> tensor conversion. +// +// Modeled on test_deepspin_dpa4_graph_ptexpt.cc (same native-spin graph +// fixture conventions) and test_deeppot_chg_spin_ptexpt.cc (the non-spin +// DeepPot twin of this charge_spin coverage). +#include + +#include +#include +#include +#include + +#include "DeepPotPTExpt.h" +#include "DeepSpin.h" +#include "DeepSpinPTExpt.h" +#include "expected_ref.h" +#include "neighbor_list.h" +#include "test_utils.h" + +// The two BUILD_PT_EXPT* macros are defined by the two PTExpt headers above +// and by nothing else. If a refactor drops either include, the #if guards +// below would silently evaluate to "skip" and this whole suite would report +// PASSED while testing nothing (that exact regression cost 8 dead cases in +// another suite). Turn that failure mode into a compile error. +#ifdef BUILD_PYTORCH +#ifndef BUILD_PT_EXPT_SPIN +#error "BUILD_PT_EXPT_SPIN undefined -- DeepSpinPTExpt.h include was dropped" +#endif +#ifndef BUILD_PT_EXPT +#error "BUILD_PT_EXPT undefined -- DeepPotPTExpt.h include was dropped" +#endif +#endif + +// Spin models need relaxed epsilon (same bound as +// test_deepspin_dpa4_graph_ptexpt.cc). +#undef EPSILON +#define EPSILON (std::is_same::value ? 1e-10 : 1e-4) + +namespace { +constexpr const char* kRefPath = + "../../tests/infer/deeppot_dpa4_spin_chgspin.expected"; +constexpr const char* kModelPath = + "../../tests/infer/deeppot_dpa4_spin_chgspin.pt2"; + +// Minimum energy separation that the two charge_spin probes must produce. +// Far above the float32 comparison bound (1e-4) so the "different energies" +// assertion means something for BOTH instantiations; the generator asserts +// the same property at 1e-6 on the fp64 reference. +constexpr double kMinChgSpinGap = 1e-3; + +// One reference section (PBC or NoPbc, default or explicit charge_spin). +template +struct SpinRefSection { + std::vector e, f, fm, tot_v, atom_v; + double tot_e = 0.; + int natoms = 0; + + void load(deepmd_test::ExpectedRef& ref, const char* section) { + e = ref.template get(section, "expected_e"); + f = ref.template get(section, "expected_f"); + fm = ref.template get(section, "expected_fm"); + tot_v = ref.template get(section, "expected_tot_v"); + atom_v = ref.template get(section, "expected_atom_v"); + natoms = static_cast(e.size()); + EXPECT_EQ(natoms * 3, static_cast(f.size())); + EXPECT_EQ(natoms * 3, static_cast(fm.size())); + EXPECT_EQ(9, static_cast(tot_v.size())); + EXPECT_EQ(natoms * 9, static_cast(atom_v.size())); + tot_e = 0.; + for (int ii = 0; ii < natoms; ++ii) { + tot_e += e[ii]; + } + } +}; + +// Compare one compute() result against a reference section. +template +void expect_matches(const SpinRefSection& ref, + double ener, + const std::vector& force, + const std::vector& force_mag, + const std::vector& virial, + double eps) { + EXPECT_EQ(force.size(), static_cast(ref.natoms * 3)); + EXPECT_EQ(force_mag.size(), static_cast(ref.natoms * 3)); + EXPECT_LT(fabs(ener - ref.tot_e), eps); + for (int ii = 0; ii < ref.natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - ref.f[ii]), eps); + EXPECT_LT(fabs(force_mag[ii] - ref.fm[ii]), eps); + } + EXPECT_FALSE(virial.empty()) << "Virial should not be empty"; + EXPECT_EQ(virial.size(), 9u); + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - ref.tot_v[ii]), eps); + } +} +} // namespace + +template +class TestInferDeepSpinDpa4ChgSpinPtExpt : public ::testing::Test { + protected: + // 6-atom system (3 Ni, spin-active; 3 O, non-magnetic) -- verbatim from + // gen_dpa4_spin_chgspin.py's _COORDS/_CELL/_SPINS/_ATYPES. + std::vector coord = {1.0, 1.0, 1.0, 3.2, 1.4, 1.1, 1.3, 1.8, 1.0, + 0.4, 1.2, 1.6, 3.6, 2.0, 1.3, 3.4, 0.7, 1.7}; + std::vector spin = {0.11, 0.05, -0.02, -0.07, 0.09, 0.03, + 0.02, -0.06, 0.08, 0.01, -0.01, 0.02, + -0.02, 0.03, -0.01, 0.015, 0.02, -0.03}; + std::vector atype = {0, 0, 0, 1, 1, 1}; + std::vector box = {6., 0., 0., 0., 6., 0., 0., 0., 6.}; + std::vector nobox = {}; + + // charge_spin is always double regardless of VALUETYPE. + // The FiLM embedding is CATEGORICAL (charge -> index charge+100, spin -> + // index spin), so both probes are integer-valued and differ in BOTH + // components: [0.0, 1.0] -> (100, 1) is the model's stored default, + // [1.0, 2.0] -> (101, 2) is the explicit runtime probe. + std::vector charge_spin_default = {0.0, 1.0}; + std::vector charge_spin_explicit = {1.0, 2.0}; + + SpinRefSection pbc_default, pbc_explicit; + SpinRefSection nopbc_default, nopbc_explicit; + + deepmd::DeepSpin dp; + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT_SPIN + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + std::ifstream model_file(kModelPath); + if (!model_file.good()) { + GTEST_SKIP() << "Skip because " << kModelPath + << " was not generated (run " + "source/tests/infer/gen_dpa4_spin_chgspin.py)."; + } + dp.init(kModelPath); + + deepmd_test::ExpectedRef ref; + ref.load(kRefPath); + pbc_default.load(ref, "pbc_default"); + pbc_explicit.load(ref, "pbc_explicit"); + nopbc_default.load(ref, "nopbc_default"); + nopbc_explicit.load(ref, "nopbc_explicit"); + + // The references themselves must be anti-vacuous: if the two sections + // carried the same energy, every "charge_spin changes the output" check + // below would pass for the wrong reason. + EXPECT_GT(fabs(pbc_explicit.tot_e - pbc_default.tot_e), kMinChgSpinGap) + << "reference sections pbc_default/pbc_explicit are degenerate; " + "regenerate with source/tests/infer/gen_dpa4_spin_chgspin.py"; + EXPECT_GT(fabs(nopbc_explicit.tot_e - nopbc_default.tot_e), kMinChgSpinGap) + << "reference sections nopbc_default/nopbc_explicit are degenerate"; + }; + + void TearDown() override {}; +}; + +TYPED_TEST_SUITE(TestInferDeepSpinDpa4ChgSpinPtExpt, ValueTypes); + +TYPED_TEST(TestInferDeepSpinDpa4ChgSpinPtExpt, dim_chg_spin) { + deepmd::DeepSpin& dp = this->dp; + // 0 here would mean the archive does not carry the charge-spin slot at all + // and every other case in this file is testing nothing. + EXPECT_EQ(dp.dim_chg_spin(), 2); +} + +TYPED_TEST(TestInferDeepSpinDpa4ChgSpinPtExpt, test_get_use_spin) { + deepmd::DeepSpin& dp = this->dp; + std::vector use_spin = dp.get_use_spin(); + EXPECT_EQ(use_spin.size(), 2); + EXPECT_TRUE(use_spin[0]); // Ni carries a magnetic moment + EXPECT_FALSE(use_spin[1]); // O does not +} + +// ============================================================================ +// Standalone (build-nlist) path -- DeepSpinPTExpt::compute, no InputNlist +// ============================================================================ + +// THE core assertion: two different runtime charge_spin vectors must give two +// different energies, and each must equal its own reference. +TYPED_TEST(TestInferDeepSpinDpa4ChgSpinPtExpt, + cpu_build_nlist_two_charge_spin) { + using VALUETYPE = TypeParam; + deepmd::DeepSpin& dp = this->dp; + + double ener_def, ener_exp; + std::vector force_def, force_mag_def, virial_def; + std::vector force_exp, force_mag_exp, virial_exp; + + dp.compute(ener_def, force_def, force_mag_def, virial_def, this->coord, + this->spin, this->atype, this->box, {}, {}, + this->charge_spin_default); + dp.compute(ener_exp, force_exp, force_mag_exp, virial_exp, this->coord, + this->spin, this->atype, this->box, {}, {}, + this->charge_spin_explicit); + + // The runtime argument reaches the model at all ... + EXPECT_GT(fabs(ener_exp - ener_def), kMinChgSpinGap) + << "charge_spin " << this->charge_spin_default[0] << "," + << this->charge_spin_default[1] << " and " + << this->charge_spin_explicit[0] << "," << this->charge_spin_explicit[1] + << " produced the same energy (" << ener_def << " vs " << ener_exp + << "): the runtime charge_spin is being ignored by the DeepSpin path."; + // ... and lands on the right values, per charge_spin. + expect_matches(this->pbc_default, ener_def, force_def, force_mag_def, + virial_def, EPSILON); + expect_matches(this->pbc_explicit, ener_exp, force_exp, force_mag_exp, + virial_exp, EPSILON); +} + +// Backward compatibility: an EMPTY charge_spin must reproduce the model's +// stored default_chg_spin -- this is what every pre-existing caller does. +TYPED_TEST(TestInferDeepSpinDpa4ChgSpinPtExpt, + cpu_build_nlist_empty_is_default) { + using VALUETYPE = TypeParam; + deepmd::DeepSpin& dp = this->dp; + + double ener_empty, ener_default_value; + std::vector f_empty, fm_empty, v_empty; + std::vector f_val, fm_val, v_val; + + // No charge_spin argument at all (the pre-existing call shape). + dp.compute(ener_empty, f_empty, fm_empty, v_empty, this->coord, this->spin, + this->atype, this->box); + expect_matches(this->pbc_default, ener_empty, f_empty, fm_empty, v_empty, + EPSILON); + + // Passing the stored default explicitly must select the same behaviour. + dp.compute(ener_default_value, f_val, fm_val, v_val, this->coord, this->spin, + this->atype, this->box, {}, {}, this->charge_spin_default); + EXPECT_LT(fabs(ener_empty - ener_default_value), EPSILON) + << "an empty charge_spin and an explicit default_chg_spin disagree"; + for (int ii = 0; ii < this->pbc_default.natoms * 3; ++ii) { + EXPECT_LT(fabs(f_empty[ii] - f_val[ii]), EPSILON); + EXPECT_LT(fabs(fm_empty[ii] - fm_val[ii]), EPSILON); + } + + // And it must NOT accidentally be the explicit-probe behaviour. + EXPECT_GT(fabs(ener_empty - this->pbc_explicit.tot_e), kMinChgSpinGap) + << "the empty-charge_spin result equals the explicit-probe reference; " + "the stored default is not being used."; +} + +TYPED_TEST(TestInferDeepSpinDpa4ChgSpinPtExpt, + cpu_build_nlist_atomic_explicit) { + using VALUETYPE = TypeParam; + deepmd::DeepSpin& dp = this->dp; + SpinRefSection& ref = this->pbc_explicit; + + double ener; + std::vector force, force_mag, virial, atom_ener, atom_vir; + dp.compute(ener, force, force_mag, virial, atom_ener, atom_vir, this->coord, + this->spin, this->atype, this->box, {}, {}, + this->charge_spin_explicit); + + expect_matches(ref, ener, force, force_mag, virial, EPSILON); + EXPECT_EQ(atom_ener.size(), static_cast(ref.natoms)); + for (int ii = 0; ii < ref.natoms; ++ii) { + EXPECT_LT(fabs(atom_ener[ii] - ref.e[ii]), EPSILON); + } + EXPECT_FALSE(atom_vir.empty()) << "Atomic virial should not be empty"; + EXPECT_EQ(atom_vir.size(), static_cast(ref.natoms * 9)); + for (int ii = 0; ii < ref.natoms * 9; ++ii) { + EXPECT_LT(fabs(atom_vir[ii] - ref.atom_v[ii]), EPSILON); + } +} + +// The size check on the runtime vector (the other branch of "charge_spin is +// non-empty"): a wrong-width charge_spin must be rejected, not silently +// truncated/padded into the model. Mirrors +// test_deeppot_chg_spin_jax.cc::rejects_invalid_input_size for DeepSpin. +TYPED_TEST(TestInferDeepSpinDpa4ChgSpinPtExpt, + rejects_invalid_charge_spin_size) { + using VALUETYPE = TypeParam; + deepmd::DeepSpin& dp = this->dp; + const std::vector invalid_charge_spin = {1.0, 2.0, 3.0}; + + double ener; + std::vector force, force_mag, virial; + EXPECT_THROW( + dp.compute(ener, force, force_mag, virial, this->coord, this->spin, + this->atype, this->box, {}, {}, invalid_charge_spin), + deepmd::deepmd_exception); +} + +// ============================================================================ +// LAMMPS path (explicit InputNlist, nghost=0) -- the SECOND +// DeepSpinPTExpt::compute overload, with its own charge_spin conversion. +// NoPBC only (no ghost atoms needed for a nghost=0 InputNlist), matching +// test_deepspin_dpa4_graph_ptexpt.cc. +// ============================================================================ + +TYPED_TEST(TestInferDeepSpinDpa4ChgSpinPtExpt, cpu_lmp_nlist_two_charge_spin) { + using VALUETYPE = TypeParam; + deepmd::DeepSpin& dp = this->dp; + const int natoms = this->nopbc_default.natoms; + + std::vector > nlist_data = { + {1, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, {0, 1, 3, 4, 5}, + {0, 1, 2, 4, 5}, {0, 1, 2, 3, 5}, {0, 1, 2, 3, 4}}; + std::vector ilist(natoms), numneigh(natoms); + std::vector firstneigh(natoms); + deepmd::InputNlist inlist(natoms, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + + double ener_def, ener_exp; + std::vector f_def, fm_def, v_def, f_exp, fm_exp, v_exp; + + // Empty charge_spin -> stored default_chg_spin. + dp.compute(ener_def, f_def, fm_def, v_def, this->coord, this->spin, + this->atype, this->nobox, 0, inlist, 0); + dp.compute(ener_exp, f_exp, fm_exp, v_exp, this->coord, this->spin, + this->atype, this->nobox, 0, inlist, 0, {}, {}, + this->charge_spin_explicit); + + EXPECT_GT(fabs(ener_exp - ener_def), kMinChgSpinGap) + << "the LAMMPS-nlist overload ignores the runtime charge_spin"; + expect_matches(this->nopbc_default, ener_def, f_def, fm_def, v_def, EPSILON); + expect_matches(this->nopbc_explicit, ener_exp, f_exp, fm_exp, v_exp, EPSILON); +} + +TYPED_TEST(TestInferDeepSpinDpa4ChgSpinPtExpt, cpu_lmp_nlist_atomic_explicit) { + using VALUETYPE = TypeParam; + deepmd::DeepSpin& dp = this->dp; + SpinRefSection& ref = this->nopbc_explicit; + const int natoms = ref.natoms; + + std::vector > nlist_data = { + {1, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, {0, 1, 3, 4, 5}, + {0, 1, 2, 4, 5}, {0, 1, 2, 3, 5}, {0, 1, 2, 3, 4}}; + std::vector ilist(natoms), numneigh(natoms); + std::vector firstneigh(natoms); + deepmd::InputNlist inlist(natoms, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + + double ener; + std::vector force, force_mag, virial, atom_ener, atom_vir; + dp.compute(ener, force, force_mag, virial, atom_ener, atom_vir, this->coord, + this->spin, this->atype, this->nobox, 0, inlist, 0, {}, {}, + this->charge_spin_explicit); + + expect_matches(ref, ener, force, force_mag, virial, EPSILON); + EXPECT_EQ(atom_ener.size(), static_cast(natoms)); + for (int ii = 0; ii < natoms; ++ii) { + EXPECT_LT(fabs(atom_ener[ii] - ref.e[ii]), EPSILON); + } + EXPECT_FALSE(atom_vir.empty()) << "Atomic virial should not be empty"; + EXPECT_EQ(atom_vir.size(), static_cast(natoms * 9)); + for (int ii = 0; ii < natoms * 9; ++ii) { + EXPECT_LT(fabs(atom_vir[ii] - ref.atom_v[ii]), EPSILON); + } +} diff --git a/source/api_cc/tests/test_deepspin_dpa4_graph_ptexpt.cc b/source/api_cc/tests/test_deepspin_dpa4_graph_ptexpt.cc new file mode 100644 index 0000000000..6e6140395e --- /dev/null +++ b/source/api_cc/tests/test_deepspin_dpa4_graph_ptexpt.cc @@ -0,0 +1,412 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Test C++ NeighborGraph inference for pt_expt (.pt2) backend with native-spin +// DPA4 (single-rank; no with-comm sibling -- has_comm_artifact=false always, +// see source/tests/infer/gen_dpa4_spin.py). Modeled on +// test_deeppot_dpa_ptexpt_spin.cc (deepspin-scheme nlist route). +#include + +#include +#include +#include +#include + +#include "DeepSpin.h" +#include "DeepSpinPTExpt.h" +#include "expected_ref.h" +#include "neighbor_list.h" +#include "test_utils.h" + +// Spin models need relaxed epsilon (same bound as +// test_deeppot_dpa_ptexpt_spin.cc). +#undef EPSILON +#define EPSILON (std::is_same::value ? 1e-10 : 1e-4) + +namespace { +constexpr const char* kRefPath = + "../../tests/infer/deeppot_dpa4_spin_graph.expected"; +constexpr const char* kModelPath = + "../../tests/infer/deeppot_dpa4_spin_graph.pt2"; +} // namespace + +// ============================================================================ +// PBC test fixture +// ============================================================================ + +template +class TestInferDeepSpinDpa4GraphPtExpt : public ::testing::Test { + protected: + // 6-atom system (3 Ni, spin-active; 3 O, non-magnetic) -- verbatim from + // gen_dpa4_spin.py's _COORDS/_CELL/_SPINS/_ATYPES. + std::vector coord = {1.0, 1.0, 1.0, 3.2, 1.4, 1.1, 1.3, 1.8, 1.0, + 0.4, 1.2, 1.6, 3.6, 2.0, 1.3, 3.4, 0.7, 1.7}; + std::vector spin = {0.11, 0.05, -0.02, -0.07, 0.09, 0.03, + 0.02, -0.06, 0.08, 0.01, -0.01, 0.02, + -0.02, 0.03, -0.01, 0.015, 0.02, -0.03}; + std::vector atype = {0, 0, 0, 1, 1, 1}; + std::vector box = {6., 0., 0., 0., 6., 0., 0., 0., 6.}; + + // Reference values generated by source/tests/infer/gen_dpa4_spin.py + std::vector expected_e; + std::vector expected_f; + std::vector expected_fm; + std::vector expected_tot_v; + std::vector expected_atom_v; + + int natoms; + double expected_tot_e; + + deepmd::DeepSpin dp; + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT_SPIN + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + std::ifstream model_file(kModelPath); + if (!model_file.good()) { + GTEST_SKIP() << "Skip because " << kModelPath + << " was not generated (run " + "source/tests/infer/gen_dpa4_spin.py)."; + } + dp.init(kModelPath); + + deepmd_test::ExpectedRef ref; + ref.load(kRefPath); + expected_e = ref.get("pbc", "expected_e"); + expected_f = ref.get("pbc", "expected_f"); + expected_fm = ref.get("pbc", "expected_fm"); + expected_tot_v = ref.get("pbc", "expected_tot_v"); + expected_atom_v = ref.get("pbc", "expected_atom_v"); + + natoms = expected_e.size(); + EXPECT_EQ(natoms * 3, expected_f.size()); + EXPECT_EQ(natoms * 3, expected_fm.size()); + EXPECT_EQ(9, expected_tot_v.size()); + EXPECT_EQ(natoms * 9, expected_atom_v.size()); + expected_tot_e = 0.; + for (int ii = 0; ii < natoms; ++ii) { + expected_tot_e += expected_e[ii]; + } + }; + + void TearDown() override {}; +}; + +TYPED_TEST_SUITE(TestInferDeepSpinDpa4GraphPtExpt, ValueTypes); + +TYPED_TEST(TestInferDeepSpinDpa4GraphPtExpt, test_get_use_spin) { + deepmd::DeepSpin& dp = this->dp; + std::vector use_spin = dp.get_use_spin(); + EXPECT_EQ(use_spin.size(), 2); + EXPECT_TRUE(use_spin[0]); // Ni carries a magnetic moment + EXPECT_FALSE(use_spin[1]); // O does not +} + +TYPED_TEST(TestInferDeepSpinDpa4GraphPtExpt, type_map) { + std::string type_map; + this->dp.get_type_map(type_map); + EXPECT_EQ(type_map, "Ni O"); +} + +// Standalone path (no InputNlist): exercises DeepSpinPTExpt::compute's +// buildGraphTensors branch (source/api_cc/src/DeepSpinPTExpt.cc). +TYPED_TEST(TestInferDeepSpinDpa4GraphPtExpt, cpu_build_nlist) { + using VALUETYPE = TypeParam; + const std::vector& coord = this->coord; + const std::vector& spin = this->spin; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_f = this->expected_f; + std::vector& expected_fm = this->expected_fm; + std::vector& expected_tot_v = this->expected_tot_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + deepmd::DeepSpin& dp = this->dp; + double ener; + std::vector force, force_mag, virial; + dp.compute(ener, force, force_mag, virial, coord, spin, atype, box); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(force_mag.size(), natoms * 3); + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), EPSILON); + } + EXPECT_FALSE(virial.empty()) << "Virial should not be empty"; + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + +TYPED_TEST(TestInferDeepSpinDpa4GraphPtExpt, cpu_build_nlist_atomic) { + using VALUETYPE = TypeParam; + const std::vector& coord = this->coord; + const std::vector& spin = this->spin; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_fm = this->expected_fm; + std::vector& expected_tot_v = this->expected_tot_v; + std::vector& expected_atom_v = this->expected_atom_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + deepmd::DeepSpin& dp = this->dp; + double ener; + std::vector force, force_mag, virial, atom_ener, atom_vir; + dp.compute(ener, force, force_mag, virial, atom_ener, atom_vir, coord, spin, + atype, box); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(force_mag.size(), natoms * 3); + EXPECT_EQ(atom_ener.size(), natoms); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), EPSILON); + } + EXPECT_FALSE(virial.empty()) << "Virial should not be empty"; + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } + for (int ii = 0; ii < natoms; ++ii) { + EXPECT_LT(fabs(atom_ener[ii] - expected_e[ii]), EPSILON); + } + EXPECT_FALSE(atom_vir.empty()) << "Atomic virial should not be empty"; + EXPECT_EQ(atom_vir.size(), natoms * 9); + for (int ii = 0; ii < natoms * 9; ++ii) { + EXPECT_LT(fabs(atom_vir[ii] - expected_atom_v[ii]), EPSILON); + } +} + +// ============================================================================ +// NoPBC test fixture +// ============================================================================ + +template +class TestInferDeepSpinDpa4GraphPtExptNopbc : public ::testing::Test { + protected: + std::vector coord = {1.0, 1.0, 1.0, 3.2, 1.4, 1.1, 1.3, 1.8, 1.0, + 0.4, 1.2, 1.6, 3.6, 2.0, 1.3, 3.4, 0.7, 1.7}; + std::vector spin = {0.11, 0.05, -0.02, -0.07, 0.09, 0.03, + 0.02, -0.06, 0.08, 0.01, -0.01, 0.02, + -0.02, 0.03, -0.01, 0.015, 0.02, -0.03}; + std::vector atype = {0, 0, 0, 1, 1, 1}; + std::vector box = {}; + + // Reference values for NoPBC from gen_dpa4_spin.py + std::vector expected_e; + std::vector expected_f; + std::vector expected_fm; + std::vector expected_tot_v; + std::vector expected_atom_v; + + int natoms; + double expected_tot_e; + + deepmd::DeepSpin dp; + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT_SPIN + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + std::ifstream model_file(kModelPath); + if (!model_file.good()) { + GTEST_SKIP() << "Skip because " << kModelPath + << " was not generated (run " + "source/tests/infer/gen_dpa4_spin.py)."; + } + dp.init(kModelPath); + + deepmd_test::ExpectedRef ref; + ref.load(kRefPath); + expected_e = ref.get("nopbc", "expected_e"); + expected_f = ref.get("nopbc", "expected_f"); + expected_fm = ref.get("nopbc", "expected_fm"); + expected_tot_v = ref.get("nopbc", "expected_tot_v"); + expected_atom_v = ref.get("nopbc", "expected_atom_v"); + + natoms = expected_e.size(); + EXPECT_EQ(natoms * 3, expected_f.size()); + EXPECT_EQ(natoms * 3, expected_fm.size()); + EXPECT_EQ(9, expected_tot_v.size()); + EXPECT_EQ(natoms * 9, expected_atom_v.size()); + expected_tot_e = 0.; + for (int ii = 0; ii < natoms; ++ii) { + expected_tot_e += expected_e[ii]; + } + }; + + void TearDown() override {}; +}; + +TYPED_TEST_SUITE(TestInferDeepSpinDpa4GraphPtExptNopbc, ValueTypes); + +TYPED_TEST(TestInferDeepSpinDpa4GraphPtExptNopbc, cpu_build_nlist) { + using VALUETYPE = TypeParam; + const std::vector& coord = this->coord; + const std::vector& spin = this->spin; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_f = this->expected_f; + std::vector& expected_fm = this->expected_fm; + std::vector& expected_tot_v = this->expected_tot_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + deepmd::DeepSpin& dp = this->dp; + double ener; + std::vector force, force_mag, virial; + dp.compute(ener, force, force_mag, virial, coord, spin, atype, box); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(force_mag.size(), natoms * 3); + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), EPSILON); + } + EXPECT_FALSE(virial.empty()) << "Virial should not be empty"; + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + +TYPED_TEST(TestInferDeepSpinDpa4GraphPtExptNopbc, cpu_build_nlist_atomic) { + using VALUETYPE = TypeParam; + const std::vector& coord = this->coord; + const std::vector& spin = this->spin; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_fm = this->expected_fm; + std::vector& expected_tot_v = this->expected_tot_v; + std::vector& expected_atom_v = this->expected_atom_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + deepmd::DeepSpin& dp = this->dp; + double ener; + std::vector force, force_mag, virial, atom_ener, atom_vir; + dp.compute(ener, force, force_mag, virial, atom_ener, atom_vir, coord, spin, + atype, box); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(force_mag.size(), natoms * 3); + EXPECT_EQ(atom_ener.size(), natoms); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), EPSILON); + } + EXPECT_FALSE(virial.empty()) << "Virial should not be empty"; + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } + for (int ii = 0; ii < natoms; ++ii) { + EXPECT_LT(fabs(atom_ener[ii] - expected_e[ii]), EPSILON); + } + EXPECT_FALSE(atom_vir.empty()) << "Atomic virial should not be empty"; + EXPECT_EQ(atom_vir.size(), natoms * 9); + for (int ii = 0; ii < natoms * 9; ++ii) { + EXPECT_LT(fabs(atom_vir[ii] - expected_atom_v[ii]), EPSILON); + } +} + +// LAMMPS path (explicit InputNlist, nghost=0): exercises +// DeepSpinPTExpt::compute's graph branch under the LAMMPS-nlist overload. +// Only exercised NoPBC (no ghost atoms needed for a nghost=0 InputNlist). +TYPED_TEST(TestInferDeepSpinDpa4GraphPtExptNopbc, cpu_lmp_nlist) { + using VALUETYPE = TypeParam; + const std::vector& coord = this->coord; + const std::vector& spin = this->spin; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_f = this->expected_f; + std::vector& expected_fm = this->expected_fm; + std::vector& expected_tot_v = this->expected_tot_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + deepmd::DeepSpin& dp = this->dp; + double ener; + std::vector force, force_mag, virial; + + std::vector > nlist_data = { + {1, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, {0, 1, 3, 4, 5}, + {0, 1, 2, 4, 5}, {0, 1, 2, 3, 5}, {0, 1, 2, 3, 4}}; + std::vector ilist(natoms), numneigh(natoms); + std::vector firstneigh(natoms); + deepmd::InputNlist inlist(natoms, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + dp.compute(ener, force, force_mag, virial, coord, spin, atype, box, 0, inlist, + 0); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(force_mag.size(), natoms * 3); + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), EPSILON); + } + EXPECT_FALSE(virial.empty()) << "Virial should not be empty"; + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + +TYPED_TEST(TestInferDeepSpinDpa4GraphPtExptNopbc, cpu_lmp_nlist_atomic) { + using VALUETYPE = TypeParam; + const std::vector& coord = this->coord; + const std::vector& spin = this->spin; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_fm = this->expected_fm; + std::vector& expected_tot_v = this->expected_tot_v; + std::vector& expected_atom_v = this->expected_atom_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + deepmd::DeepSpin& dp = this->dp; + double ener; + std::vector force, force_mag, virial, atom_ener, atom_vir; + + std::vector > nlist_data = { + {1, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, {0, 1, 3, 4, 5}, + {0, 1, 2, 4, 5}, {0, 1, 2, 3, 5}, {0, 1, 2, 3, 4}}; + std::vector ilist(natoms), numneigh(natoms); + std::vector firstneigh(natoms); + deepmd::InputNlist inlist(natoms, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + dp.compute(ener, force, force_mag, virial, atom_ener, atom_vir, coord, spin, + atype, box, 0, inlist, 0); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(force_mag.size(), natoms * 3); + EXPECT_EQ(atom_ener.size(), natoms); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), EPSILON); + } + EXPECT_FALSE(virial.empty()) << "Virial should not be empty"; + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } + for (int ii = 0; ii < natoms; ++ii) { + EXPECT_LT(fabs(atom_ener[ii] - expected_e[ii]), EPSILON); + } + EXPECT_FALSE(atom_vir.empty()) << "Atomic virial should not be empty"; + EXPECT_EQ(atom_vir.size(), natoms * 9); + for (int ii = 0; ii < natoms * 9; ++ii) { + EXPECT_LT(fabs(atom_vir[ii] - expected_atom_v[ii]), EPSILON); + } +} diff --git a/source/api_cc/tests/test_deepspin_dpa4_pairexcl_ptexpt.cc b/source/api_cc/tests/test_deepspin_dpa4_pairexcl_ptexpt.cc new file mode 100644 index 0000000000..1c16761c7f --- /dev/null +++ b/source/api_cc/tests/test_deepspin_dpa4_pairexcl_ptexpt.cc @@ -0,0 +1,234 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// C++ model-level pair-exclusion seam for the native-spin graph route +// (DeepSpinPTExpt). Twin of test_deeppot_dpa1_pairexcl_ptexpt.cc, which +// covers the non-spin DeepPotPTExpt seam. +// +// Model-level exclusion is a BUILD-time transform owned by the neighbor-graph +// construction (decision #18/A4): the exported .pt2 lower consumes a +// pre-excluded ``edge_mask`` and never re-applies it, so every external feeder +// must fold it in. ``DeepSpinPTExpt`` therefore parses ``pair_exclude_types`` +// from metadata once in ``init`` and calls ``applyPairExclusion`` in BOTH +// graph branches (cached-nlist and standalone) -- these tests are what proves +// those calls are load-bearing. +// +// The fixture (source/tests/infer/gen_dpa4_spin.py) is deliberately +// anti-vacuous: ``deeppot_dpa4_spin_pairexcl.pt2`` carries +// ``pair_exclude_types=[[0, 1]]`` at the MODEL level with the descriptor's own +// ``exclude_types`` left EMPTY, and shares byte-identical weights with the +// no-exclusion baseline ``deeppot_dpa4_spin_graph.pt2``. Nothing inside the +// compiled artifact reproduces the mask, so a dead C++ seam changes the +// numbers. (The ``type="dpa4"`` model alias copies model-level pairs into +// ``descriptor.exclude_types``, which WOULD bake an equivalent mask into the +// artifact and hide exactly this bug -- hence the generic +// ``type="standard"``-style config in the generator.) +// +// Two assertions per ingestion branch: +// 1. C++ == the Python DeepEval reference for the SAME archive (1e-10), +// i.e. the seam applies the exclusion the way Python does; +// 2. excluded != baseline, i.e. the exclusion is genuinely active. +#include + +#include +#include +#include +#include + +#include "DeepSpin.h" +#include "DeepSpinPTExpt.h" +#include "expected_ref.h" +#include "neighbor_list.h" +#include "test_utils.h" + +// Spin models need relaxed epsilon (same bound as +// test_deepspin_dpa4_graph_ptexpt.cc). +#undef EPSILON +#define EPSILON (std::is_same::value ? 1e-10 : 1e-4) + +namespace { +constexpr const char* kExclModel = + "../../tests/infer/deeppot_dpa4_spin_pairexcl.pt2"; +constexpr const char* kExclRef = + "../../tests/infer/deeppot_dpa4_spin_pairexcl.expected"; +constexpr const char* kBaseModel = + "../../tests/infer/deeppot_dpa4_spin_graph.pt2"; +} // namespace + +template +class TestInferDeepSpinDpa4PairExclPtExpt : public ::testing::Test { + protected: + // 6-atom system (3 Ni spin-active, 3 O non-magnetic) -- verbatim from + // gen_dpa4_spin.py's _COORDS/_CELL/_SPINS/_ATYPES. With + // pair_exclude_types=[[0, 1]] every Ni-O edge is dropped, so only the + // Ni-Ni and O-O interactions survive. + std::vector coord = {1.0, 1.0, 1.0, 3.2, 1.4, 1.1, 1.3, 1.8, 1.0, + 0.4, 1.2, 1.6, 3.6, 2.0, 1.3, 3.4, 0.7, 1.7}; + std::vector spin = {0.11, 0.05, -0.02, -0.07, 0.09, 0.03, + 0.02, -0.06, 0.08, 0.01, -0.01, 0.02, + -0.02, 0.03, -0.01, 0.015, 0.02, -0.03}; + std::vector atype = {0, 0, 0, 1, 1, 1}; + std::vector box = {6., 0., 0., 0., 6., 0., 0., 0., 6.}; + + std::vector expected_e; + std::vector expected_f; + std::vector expected_fm; + std::vector expected_tot_v; + // NoPBC twin: an explicit nghost=0 InputNlist carries no periodic images, + // so the LAMMPS-nlist case below is the gas-phase system and must be + // compared against the gas-phase reference (same convention as + // test_deepspin_dpa4_graph_ptexpt.cc, which runs cpu_lmp_nlist only in its + // Nopbc fixture). + std::vector expected_e_nopbc; + std::vector expected_f_nopbc; + std::vector expected_fm_nopbc; + + int natoms; + double expected_tot_e; + double expected_tot_e_nopbc; + + deepmd::DeepSpin dp_excl; + deepmd::DeepSpin dp_base; + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT_SPIN + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + std::ifstream excl_file(kExclModel); + std::ifstream base_file(kBaseModel); + if (!excl_file.good() || !base_file.good()) { + GTEST_SKIP() << "Skip because the native-spin DPA4 fixtures were not " + "generated (run source/tests/infer/gen_dpa4_spin.py)."; + } + dp_excl.init(kExclModel); + dp_base.init(kBaseModel); + + deepmd_test::ExpectedRef ref; + ref.load(kExclRef); + expected_e = ref.get("pbc", "expected_e"); + expected_f = ref.get("pbc", "expected_f"); + expected_fm = ref.get("pbc", "expected_fm"); + expected_tot_v = ref.get("pbc", "expected_tot_v"); + expected_e_nopbc = ref.get("nopbc", "expected_e"); + expected_f_nopbc = ref.get("nopbc", "expected_f"); + expected_fm_nopbc = ref.get("nopbc", "expected_fm"); + + natoms = expected_e.size(); + EXPECT_EQ(natoms * 3, expected_f.size()); + EXPECT_EQ(natoms * 3, expected_fm.size()); + EXPECT_EQ(9, expected_tot_v.size()); + expected_tot_e = 0.; + expected_tot_e_nopbc = 0.; + for (int ii = 0; ii < natoms; ++ii) { + expected_tot_e += expected_e[ii]; + expected_tot_e_nopbc += expected_e_nopbc[ii]; + } + }; + + void TearDown() override {}; +}; + +TYPED_TEST_SUITE(TestInferDeepSpinDpa4PairExclPtExpt, ValueTypes); + +// Standalone branch: DeepSpinPTExpt::compute -> buildGraphTensors -> +// applyPairExclusion. +TYPED_TEST(TestInferDeepSpinDpa4PairExclPtExpt, cpu_build_nlist) { + using VALUETYPE = TypeParam; + const std::vector& coord = this->coord; + const std::vector& spin = this->spin; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_f = this->expected_f; + std::vector& expected_fm = this->expected_fm; + std::vector& expected_tot_v = this->expected_tot_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + + double ener; + std::vector force, force_mag, virial; + this->dp_excl.compute(ener, force, force_mag, virial, coord, spin, atype, + box); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(force_mag.size(), natoms * 3); + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON) + << "model-level pair_exclude_types is dropped on the DeepSpinPTExpt " + "standalone graph seam"; + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), EPSILON); + } + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + +// Cached-nlist (LAMMPS) branch: DeepSpinPTExpt::compute_inner -> +// compactEdgeTensors -> applyPairExclusion. +TYPED_TEST(TestInferDeepSpinDpa4PairExclPtExpt, cpu_lmp_nlist) { + using VALUETYPE = TypeParam; + const std::vector& coord = this->coord; + const std::vector& spin = this->spin; + std::vector& atype = this->atype; + std::vector& expected_f = this->expected_f_nopbc; + std::vector& expected_fm = this->expected_fm_nopbc; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e_nopbc; + // An nghost=0 InputNlist carries no periodic images, so this is the + // gas-phase system: empty box + the NoPBC reference. + std::vector box = {}; + + double ener; + std::vector force, force_mag, virial; + // All-pairs nlist: the model's own rcut cut and the exclusion must both be + // applied at the ingestion seam. + std::vector > nlist_data = { + {1, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, {0, 1, 3, 4, 5}, + {0, 1, 2, 4, 5}, {0, 1, 2, 3, 5}, {0, 1, 2, 3, 4}}; + std::vector ilist(natoms), numneigh(natoms); + std::vector firstneigh(natoms); + deepmd::InputNlist inlist(natoms, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + this->dp_excl.compute(ener, force, force_mag, virial, coord, spin, atype, box, + 0, inlist, 0); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON) + << "model-level pair_exclude_types is dropped on the DeepSpinPTExpt " + "cached-nlist graph seam"; + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), EPSILON); + } +} + +// Anti-vacuity: the two archives share weights and differ ONLY by the +// model-level exclusion, so equal predictions would mean the exclusion never +// reached the graph build and both comparisons above would pass for the wrong +// reason. +TYPED_TEST(TestInferDeepSpinDpa4PairExclPtExpt, + excluded_differs_from_baseline) { + using VALUETYPE = TypeParam; + const std::vector& coord = this->coord; + const std::vector& spin = this->spin; + std::vector& atype = this->atype; + std::vector& box = this->box; + int& natoms = this->natoms; + + double ener_excl, ener_base; + std::vector f_excl, fm_excl, v_excl; + std::vector f_base, fm_base, v_base; + this->dp_excl.compute(ener_excl, f_excl, fm_excl, v_excl, coord, spin, atype, + box); + this->dp_base.compute(ener_base, f_base, fm_base, v_base, coord, spin, atype, + box); + + EXPECT_GT(fabs(ener_excl - ener_base), 1e-6) + << "excluding every Ni-O pair left the energy unchanged; the fixture or " + "the exclusion seam is vacuous"; + double max_df = 0.; + for (int ii = 0; ii < natoms * 3; ++ii) { + max_df = + std::max(max_df, static_cast(fabs(f_excl[ii] - f_base[ii]))); + } + EXPECT_GT(max_df, 1e-6) + << "excluding every Ni-O pair left the forces unchanged"; +} diff --git a/source/api_cc/tests/test_deepspin_dpa4_zbl_ptexpt.cc b/source/api_cc/tests/test_deepspin_dpa4_zbl_ptexpt.cc new file mode 100644 index 0000000000..a60dc8863e --- /dev/null +++ b/source/api_cc/tests/test_deepspin_dpa4_zbl_ptexpt.cc @@ -0,0 +1,509 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// C++ NeighborGraph inference for a NATIVE-SPIN DPA4 that is ALSO bridged +// with the analytical ZBL term (pt_expt .pt2, graph lower). +// +// The combination is not the conjunction of the two paths already covered by +// test_deepspin_dpa4_graph_ptexpt.cc (native spin) and +// test_deeppot_dpa4_zbl_ptexpt.cc (bridging): here ``spin`` must reach a +// COMPOSITION -- LinearEnergyModel over [learned DPA4, +// InterPotentialAtomicModel] -- whose learned child consumes it and whose +// analytical child accepts and ignores it, and the archive must still declare +// ``is_spin`` so this DeepSpin path (not DeepPot) is the one that runs. +// Before this test the combination had NO coverage below the pt_expt Python +// layer, and its only export-seam test skips when CI=true. +// +// Single-rank only by construction: bridging enables the descriptor's Source +// Freeze Propagation Gate, whose per-node eta folds a node's full +// outgoing-edge set, and edges exist only for owned centres -- so +// source/tests/infer/gen_dpa4_spin_zbl.py asserts has_comm_artifact=false AND +// that no nested forward_lower_with_comm.pt2 exists. Multi-rank is therefore +// out of scope here; under mpirun the C++ dispatch fails fast (asserted from +// the LAMMPS side, source/lmp/tests/test_lammps_dpa4_zbl_pt2.py). +// +// The references come from the Python DeepEval of the SAME archive +// (gen_dpa4_spin_zbl.py), which the generator in turn holds to its eager +// dpmodel, so a match here validates the whole C++ chain: metadata parse, +// graph + spin ingestion, and the compiled math of the composition. +#include + +#include +#include +#include +#include +#include +#include + +#include "DeepSpin.h" +// Defines BUILD_PT_EXPT_SPIN (via its __has_include probe for the inductor +// headers), which the SetUp guards below read. +#include "DeepSpinPTExpt.h" +// Defines BUILD_PT_EXPT, same probe. Without these includes the guards see +// the macros undefined and every case GTEST_SKIPs with "PyTorch support is +// not enabled" -- silently, which is how the ZBL suite first ran as 8 skips +// that looked like passes. The #errors below make that failure mode +// impossible to reintroduce: a missing include is now a build error, not a +// silent skip. +#include "DeepPotPTExpt.h" + +#if defined(BUILD_PYTORCH) && !defined(BUILD_PT_EXPT) +#error "BUILD_PT_EXPT undefined: DeepPotPTExpt.h must be included." +#endif +#if defined(BUILD_PYTORCH) && !defined(BUILD_PT_EXPT_SPIN) +#error "BUILD_PT_EXPT_SPIN undefined: DeepSpinPTExpt.h must be included." +#endif + +#include "expected_ref.h" +#include "neighbor_list.h" +#include "test_utils.h" + +namespace { +constexpr const char* kModelPath = + "../../tests/infer/deeppot_dpa4_spin_zbl_graph.pt2"; +constexpr const char* kRefPath = + "../../tests/infer/deeppot_dpa4_spin_zbl_graph.expected"; + +// Magnitude-scaled bound, same reasoning as test_deeppot_dpa4_zbl_ptexpt.cc's +// ``zbl_tol``: the analytical ZBL term on the fixture's 0.9 A Ni-Ni pair +// makes the forces ~1.4e3 and the virial ~1.3e3 (deliberately, so the term +// cannot be mistaken for noise), and ONE fp32 ULP at that magnitude is +// 2^-13 = 1.22e-4. The suite's flat 1e-4 float bound therefore asks for +// sub-ULP agreement, which no fp32 result can satisfy. 5e-7 relative is +// ~4 fp32 ULP; the 1e-4 floor keeps near-zero components (notably force_mag, +// whose largest entry is ~8e-3) at the suite's usual float bound. +// +// fp64 keeps a strict absolute 1e-10 -- the same bound the sibling native-spin +// graph suite uses, and still ~400x tighter than one fp64 ULP at this +// fixture's force magnitude (2.3e-13). +template +inline double zbl_spin_tol(double expected) { + if (std::is_same::value) { + return 1e-10; + } + return 1e-4 + 5e-7 * fabs(expected); +} +} // namespace + +// ============================================================================ +// PBC test fixture +// ============================================================================ + +template +class TestInferDeepSpinDpa4ZblPtExpt : public ::testing::Test { + protected: + // 6-atom system (3 Ni, spin-active; 3 O, non-magnetic) -- verbatim from + // gen_dpa4_spin_zbl.py's _COORDS/_CELL/_SPINS/_ATYPES. Atoms 0 and 1 sit + // 0.9 A apart, inside bridging_r_outer, and are BOTH spin-carrying Ni, so + // the analytical and spin channels act on the same atoms. + std::vector coord = {1.0, 1.0, 1.0, 1.9, 1.0, 1.0, 1.3, 1.8, 1.0, + 0.4, 1.2, 1.6, 3.6, 2.0, 1.3, 3.4, 0.7, 1.7}; + std::vector spin = {0.11, 0.05, -0.02, -0.07, 0.09, 0.03, + 0.02, -0.06, 0.08, 0.01, -0.01, 0.02, + -0.02, 0.03, -0.01, 0.015, 0.02, -0.03}; + std::vector atype = {0, 0, 0, 1, 1, 1}; + std::vector box = {6., 0., 0., 0., 6., 0., 0., 0., 6.}; + + std::vector expected_e; + std::vector expected_f; + std::vector expected_fm; + std::vector expected_tot_v; + std::vector expected_atom_v; + + int natoms; + double expected_tot_e; + + deepmd::DeepSpin dp; + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT_SPIN + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + std::ifstream model_file(kModelPath); + if (!model_file.good()) { + GTEST_SKIP() << "Skip because " << kModelPath + << " was not generated (run " + "source/tests/infer/gen_dpa4_spin_zbl.py)."; + } + dp.init(kModelPath); + + deepmd_test::ExpectedRef ref; + ref.load(kRefPath); + expected_e = ref.get("pbc", "expected_e"); + expected_f = ref.get("pbc", "expected_f"); + expected_fm = ref.get("pbc", "expected_fm"); + expected_tot_v = ref.get("pbc", "expected_tot_v"); + expected_atom_v = ref.get("pbc", "expected_atom_v"); + + natoms = expected_e.size(); + EXPECT_EQ(natoms * 3, expected_f.size()); + EXPECT_EQ(natoms * 3, expected_fm.size()); + EXPECT_EQ(9, expected_tot_v.size()); + EXPECT_EQ(natoms * 9, expected_atom_v.size()); + expected_tot_e = 0.; + for (int ii = 0; ii < natoms; ++ii) { + expected_tot_e += expected_e[ii]; + } + }; + + void TearDown() override {}; +}; + +TYPED_TEST_SUITE(TestInferDeepSpinDpa4ZblPtExpt, ValueTypes); + +TYPED_TEST(TestInferDeepSpinDpa4ZblPtExpt, test_get_use_spin) { + deepmd::DeepSpin& dp = this->dp; + std::vector use_spin = dp.get_use_spin(); + EXPECT_EQ(use_spin.size(), 2); + EXPECT_TRUE(use_spin[0]); // Ni carries a magnetic moment + EXPECT_FALSE(use_spin[1]); // O does not +} + +TYPED_TEST(TestInferDeepSpinDpa4ZblPtExpt, type_map) { + std::string type_map; + this->dp.get_type_map(type_map); + EXPECT_EQ(type_map, "Ni O"); +} + +// Standalone path (no InputNlist): exercises DeepSpinPTExpt::compute's +// buildGraphTensors branch on a linear composition. +TYPED_TEST(TestInferDeepSpinDpa4ZblPtExpt, cpu_build_nlist) { + using VALUETYPE = TypeParam; + const std::vector& coord = this->coord; + const std::vector& spin = this->spin; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_f = this->expected_f; + std::vector& expected_fm = this->expected_fm; + std::vector& expected_tot_v = this->expected_tot_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + deepmd::DeepSpin& dp = this->dp; + double ener; + std::vector force, force_mag, virial; + dp.compute(ener, force, force_mag, virial, coord, spin, atype, box); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(force_mag.size(), natoms * 3); + + // Anti-vacuity, on the values this run actually produced (not only on the + // reference): the close Ni-Ni pair must drive a large ZBL repulsion, and + // the jittered learned half must drive a nonzero force_mag on the Ni atoms + // while the model's own type gate keeps the O rows exactly zero. + double fmax = 0.; + for (int ii = 0; ii < natoms * 3; ++ii) { + fmax = std::max(fmax, static_cast(fabs(force[ii]))); + } + EXPECT_GT(fmax, 1e-3) << "forces are trivially small; fixture is vacuous"; + double fm_spin_max = 0.; + for (int ii = 0; ii < natoms; ++ii) { + for (int dd = 0; dd < 3; ++dd) { + double v = fabs(force_mag[ii * 3 + dd]); + if (atype[ii] == 0) { + fm_spin_max = std::max(fm_spin_max, v); + } else { + EXPECT_EQ(force_mag[ii * 3 + dd], static_cast(0)) + << "force_mag must be EXACTLY zero on non-spin (O) atom " << ii; + } + } + } + EXPECT_GT(fm_spin_max, 1e-6) + << "force_mag is trivially small on the spin-active atoms"; + + EXPECT_LT(fabs(ener - expected_tot_e), + zbl_spin_tol(expected_tot_e)); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), + zbl_spin_tol(expected_f[ii])); + EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), + zbl_spin_tol(expected_fm[ii])); + } + EXPECT_FALSE(virial.empty()) << "Virial should not be empty"; + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), + zbl_spin_tol(expected_tot_v[ii])); + } +} + +TYPED_TEST(TestInferDeepSpinDpa4ZblPtExpt, cpu_build_nlist_atomic) { + using VALUETYPE = TypeParam; + const std::vector& coord = this->coord; + const std::vector& spin = this->spin; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_fm = this->expected_fm; + std::vector& expected_tot_v = this->expected_tot_v; + std::vector& expected_atom_v = this->expected_atom_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + deepmd::DeepSpin& dp = this->dp; + double ener; + std::vector force, force_mag, virial, atom_ener, atom_vir; + dp.compute(ener, force, force_mag, virial, atom_ener, atom_vir, coord, spin, + atype, box); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(force_mag.size(), natoms * 3); + EXPECT_EQ(atom_ener.size(), natoms); + + EXPECT_LT(fabs(ener - expected_tot_e), + zbl_spin_tol(expected_tot_e)); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), + zbl_spin_tol(expected_f[ii])); + EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), + zbl_spin_tol(expected_fm[ii])); + } + EXPECT_FALSE(virial.empty()) << "Virial should not be empty"; + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), + zbl_spin_tol(expected_tot_v[ii])); + } + for (int ii = 0; ii < natoms; ++ii) { + EXPECT_LT(fabs(atom_ener[ii] - expected_e[ii]), + zbl_spin_tol(expected_e[ii])); + } + EXPECT_FALSE(atom_vir.empty()) << "Atomic virial should not be empty"; + EXPECT_EQ(atom_vir.size(), natoms * 9); + for (int ii = 0; ii < natoms * 9; ++ii) { + EXPECT_LT(fabs(atom_vir[ii] - expected_atom_v[ii]), + zbl_spin_tol(expected_atom_v[ii])); + } +} + +// ============================================================================ +// NoPBC test fixture +// ============================================================================ + +template +class TestInferDeepSpinDpa4ZblPtExptNopbc : public ::testing::Test { + protected: + std::vector coord = {1.0, 1.0, 1.0, 1.9, 1.0, 1.0, 1.3, 1.8, 1.0, + 0.4, 1.2, 1.6, 3.6, 2.0, 1.3, 3.4, 0.7, 1.7}; + std::vector spin = {0.11, 0.05, -0.02, -0.07, 0.09, 0.03, + 0.02, -0.06, 0.08, 0.01, -0.01, 0.02, + -0.02, 0.03, -0.01, 0.015, 0.02, -0.03}; + std::vector atype = {0, 0, 0, 1, 1, 1}; + std::vector box = {}; + + std::vector expected_e; + std::vector expected_f; + std::vector expected_fm; + std::vector expected_tot_v; + std::vector expected_atom_v; + + int natoms; + double expected_tot_e; + + deepmd::DeepSpin dp; + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT_SPIN + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + std::ifstream model_file(kModelPath); + if (!model_file.good()) { + GTEST_SKIP() << "Skip because " << kModelPath + << " was not generated (run " + "source/tests/infer/gen_dpa4_spin_zbl.py)."; + } + dp.init(kModelPath); + + deepmd_test::ExpectedRef ref; + ref.load(kRefPath); + expected_e = ref.get("nopbc", "expected_e"); + expected_f = ref.get("nopbc", "expected_f"); + expected_fm = ref.get("nopbc", "expected_fm"); + expected_tot_v = ref.get("nopbc", "expected_tot_v"); + expected_atom_v = ref.get("nopbc", "expected_atom_v"); + + natoms = expected_e.size(); + EXPECT_EQ(natoms * 3, expected_f.size()); + EXPECT_EQ(natoms * 3, expected_fm.size()); + EXPECT_EQ(9, expected_tot_v.size()); + EXPECT_EQ(natoms * 9, expected_atom_v.size()); + expected_tot_e = 0.; + for (int ii = 0; ii < natoms; ++ii) { + expected_tot_e += expected_e[ii]; + } + }; + + void TearDown() override {}; +}; + +TYPED_TEST_SUITE(TestInferDeepSpinDpa4ZblPtExptNopbc, ValueTypes); + +TYPED_TEST(TestInferDeepSpinDpa4ZblPtExptNopbc, cpu_build_nlist) { + using VALUETYPE = TypeParam; + const std::vector& coord = this->coord; + const std::vector& spin = this->spin; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_f = this->expected_f; + std::vector& expected_fm = this->expected_fm; + std::vector& expected_tot_v = this->expected_tot_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + deepmd::DeepSpin& dp = this->dp; + double ener; + std::vector force, force_mag, virial; + dp.compute(ener, force, force_mag, virial, coord, spin, atype, box); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(force_mag.size(), natoms * 3); + EXPECT_LT(fabs(ener - expected_tot_e), + zbl_spin_tol(expected_tot_e)); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), + zbl_spin_tol(expected_f[ii])); + EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), + zbl_spin_tol(expected_fm[ii])); + } + EXPECT_FALSE(virial.empty()) << "Virial should not be empty"; + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), + zbl_spin_tol(expected_tot_v[ii])); + } +} + +TYPED_TEST(TestInferDeepSpinDpa4ZblPtExptNopbc, cpu_build_nlist_atomic) { + using VALUETYPE = TypeParam; + const std::vector& coord = this->coord; + const std::vector& spin = this->spin; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_fm = this->expected_fm; + std::vector& expected_tot_v = this->expected_tot_v; + std::vector& expected_atom_v = this->expected_atom_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + deepmd::DeepSpin& dp = this->dp; + double ener; + std::vector force, force_mag, virial, atom_ener, atom_vir; + dp.compute(ener, force, force_mag, virial, atom_ener, atom_vir, coord, spin, + atype, box); + + EXPECT_EQ(atom_ener.size(), natoms); + EXPECT_LT(fabs(ener - expected_tot_e), + zbl_spin_tol(expected_tot_e)); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), + zbl_spin_tol(expected_f[ii])); + EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), + zbl_spin_tol(expected_fm[ii])); + } + for (int ii = 0; ii < natoms; ++ii) { + EXPECT_LT(fabs(atom_ener[ii] - expected_e[ii]), + zbl_spin_tol(expected_e[ii])); + } + EXPECT_FALSE(virial.empty()) << "Virial should not be empty"; + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), + zbl_spin_tol(expected_tot_v[ii])); + } + EXPECT_FALSE(atom_vir.empty()) << "Atomic virial should not be empty"; + EXPECT_EQ(atom_vir.size(), natoms * 9); + for (int ii = 0; ii < natoms * 9; ++ii) { + EXPECT_LT(fabs(atom_vir[ii] - expected_atom_v[ii]), + zbl_spin_tol(expected_atom_v[ii])); + } +} + +// LAMMPS path (explicit InputNlist, nghost=0): exercises DeepSpinPTExpt's +// graph branch under the LAMMPS-nlist overload. Gas-phase system, so it is +// compared against the NoPBC reference. +TYPED_TEST(TestInferDeepSpinDpa4ZblPtExptNopbc, cpu_lmp_nlist) { + using VALUETYPE = TypeParam; + const std::vector& coord = this->coord; + const std::vector& spin = this->spin; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_f = this->expected_f; + std::vector& expected_fm = this->expected_fm; + std::vector& expected_tot_v = this->expected_tot_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + deepmd::DeepSpin& dp = this->dp; + double ener; + std::vector force, force_mag, virial; + + std::vector > nlist_data = { + {1, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, {0, 1, 3, 4, 5}, + {0, 1, 2, 4, 5}, {0, 1, 2, 3, 5}, {0, 1, 2, 3, 4}}; + std::vector ilist(natoms), numneigh(natoms); + std::vector firstneigh(natoms); + deepmd::InputNlist inlist(natoms, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + dp.compute(ener, force, force_mag, virial, coord, spin, atype, box, 0, inlist, + 0); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(force_mag.size(), natoms * 3); + EXPECT_LT(fabs(ener - expected_tot_e), + zbl_spin_tol(expected_tot_e)); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), + zbl_spin_tol(expected_f[ii])); + EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), + zbl_spin_tol(expected_fm[ii])); + } + EXPECT_FALSE(virial.empty()) << "Virial should not be empty"; + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), + zbl_spin_tol(expected_tot_v[ii])); + } +} + +TYPED_TEST(TestInferDeepSpinDpa4ZblPtExptNopbc, cpu_lmp_nlist_atomic) { + using VALUETYPE = TypeParam; + const std::vector& coord = this->coord; + const std::vector& spin = this->spin; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_fm = this->expected_fm; + std::vector& expected_atom_v = this->expected_atom_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + deepmd::DeepSpin& dp = this->dp; + double ener; + std::vector force, force_mag, virial, atom_ener, atom_vir; + + std::vector > nlist_data = { + {1, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, {0, 1, 3, 4, 5}, + {0, 1, 2, 4, 5}, {0, 1, 2, 3, 5}, {0, 1, 2, 3, 4}}; + std::vector ilist(natoms), numneigh(natoms); + std::vector firstneigh(natoms); + deepmd::InputNlist inlist(natoms, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + dp.compute(ener, force, force_mag, virial, atom_ener, atom_vir, coord, spin, + atype, box, 0, inlist, 0); + + EXPECT_EQ(atom_ener.size(), natoms); + EXPECT_LT(fabs(ener - expected_tot_e), + zbl_spin_tol(expected_tot_e)); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), + zbl_spin_tol(expected_f[ii])); + EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), + zbl_spin_tol(expected_fm[ii])); + } + for (int ii = 0; ii < natoms; ++ii) { + EXPECT_LT(fabs(atom_ener[ii] - expected_e[ii]), + zbl_spin_tol(expected_e[ii])); + } + EXPECT_FALSE(atom_vir.empty()) << "Atomic virial should not be empty"; + EXPECT_EQ(atom_vir.size(), natoms * 9); + for (int ii = 0; ii < natoms * 9; ++ii) { + EXPECT_LT(fabs(atom_vir[ii] - expected_atom_v[ii]), + zbl_spin_tol(expected_atom_v[ii])); + } +} diff --git a/source/install/test_cc_local.sh b/source/install/test_cc_local.sh index 118d480d19..d31a883607 100755 --- a/source/install/test_cc_local.sh +++ b/source/install/test_cc_local.sh @@ -142,6 +142,31 @@ else: PID8=$! wait $PID7 wait $PID8 + + # Native-spin DPA4 graph archives (baseline + model-level pair + # exclusion). Without this the whole native-spin graph C++ suite + # GTEST_SKIPs on the missing fixture, which is how a dead + # applyPairExclusion seam in DeepSpinPTExpt went unnoticed. + env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4_spin.py & + PID11=$! + # DPA4 + analytical ZBL bridging: a linear COMPOSITION on the graph + # lower, which no other C++ fixture covers. + env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4_zbl.py & + PID12=$! + # Native-spin DPA4 + charge-spin FiLM: the ONLY fixture with both + # is_spin=true and dim_chg_spin>0, i.e. the only one on which the + # runtime charge_spin argument of DeepSpin::compute is not inert. + env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4_spin_chgspin.py & + PID13=$! + # Native-spin DPA4 + ZBL bridging COMBINED: spin reaching a linear + # composition, and the only fixture pinning that a bridged model + # exports NO with-comm artifact (single-rank only by construction). + env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4_spin_zbl.py & + PID14=$! + wait $PID11 + wait $PID12 + wait $PID13 + wait $PID14 fi fi diff --git a/source/lmp/pair_deepspin.cpp b/source/lmp/pair_deepspin.cpp index aba67fae6c..2f6abe8c3e 100644 --- a/source/lmp/pair_deepspin.cpp +++ b/source/lmp/pair_deepspin.cpp @@ -260,8 +260,8 @@ void PairDeepSpin::compute(int eflag, int vflag) { if (!(eflag_atom || cvflag_atom)) { try { deep_spin.compute(dener, dforce, dforce_mag, dvirial, dcoord, dspin, - dtype, dbox, nghost, lmp_list, ago, fparam, - daparam); + dtype, dbox, nghost, lmp_list, ago, fparam, daparam, + charge_spin); } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } @@ -273,7 +273,7 @@ void PairDeepSpin::compute(int eflag, int vflag) { try { deep_spin.compute(dener, dforce, dforce_mag, dvirial, deatom, dvatom, dcoord, dspin, dtype, dbox, nghost, lmp_list, ago, - fparam, daparam); + fparam, daparam, charge_spin); } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } @@ -316,9 +316,9 @@ void PairDeepSpin::compute(int eflag, int vflag) { vector> all_atom_virial; if (!(eflag_atom || cvflag_atom)) { try { - deep_spin_model_devi.compute(all_energy, all_force, all_force_mag, - all_virial, dcoord, dspin, dtype, dbox, - nghost, lmp_list, ago, fparam, daparam); + deep_spin_model_devi.compute( + all_energy, all_force, all_force_mag, all_virial, dcoord, dspin, + dtype, dbox, nghost, lmp_list, ago, fparam, daparam, charge_spin); } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } @@ -327,7 +327,7 @@ void PairDeepSpin::compute(int eflag, int vflag) { deep_spin_model_devi.compute( all_energy, all_force, all_force_mag, all_virial, all_atom_energy, all_atom_virial, dcoord, dspin, dtype, dbox, nghost, lmp_list, - ago, fparam, daparam); + ago, fparam, daparam, charge_spin); } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } @@ -506,7 +506,8 @@ void PairDeepSpin::compute(int eflag, int vflag) { if (numb_models == 1) { try { deep_spin.compute(dener, dforce, dforce_mag, dvirial, dcoord, dspin, - dtype, dbox); + dtype, dbox, vector(), vector(), + charge_spin); } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } @@ -552,6 +553,7 @@ static bool is_key(const string& input) { keys.push_back("aparam"); keys.push_back("fparam_from_compute"); keys.push_back("aparam_from_compute"); + keys.push_back("charge_spin"); keys.push_back("ttm"); keys.push_back("atomic"); keys.push_back("relative"); @@ -595,6 +597,7 @@ void PairDeepSpin::settings(int narg, char** arg) { numb_types_spin = deep_spin.numb_types_spin(); dim_fparam = deep_spin.dim_fparam(); dim_aparam = deep_spin.dim_aparam(); + dim_chg_spin = deep_spin.dim_chg_spin(); } else { try { deep_spin.init(arg[0], get_node_rank(), get_file_content(arg[0])); @@ -608,11 +611,13 @@ void PairDeepSpin::settings(int narg, char** arg) { numb_types_spin = deep_spin_model_devi.numb_types_spin(); dim_fparam = deep_spin_model_devi.dim_fparam(); dim_aparam = deep_spin_model_devi.dim_aparam(); + dim_chg_spin = deep_spin_model_devi.dim_chg_spin(); assert(cutoff == deep_spin.cutoff() * dist_unit_cvt_factor); assert(numb_types == deep_spin.numb_types()); assert(numb_types_spin == deep_spin.numb_types_spin()); assert(dim_fparam == deep_spin.dim_fparam()); assert(dim_aparam == deep_spin.dim_aparam()); + assert(dim_chg_spin == deep_spin.dim_chg_spin()); } out_freq = 100; @@ -622,6 +627,7 @@ void PairDeepSpin::settings(int narg, char** arg) { eps = 0.; fparam.clear(); aparam.clear(); + charge_spin.clear(); while (iarg < narg) { if (!is_key(arg[iarg])) { error->all(FLERR, @@ -661,6 +667,17 @@ void PairDeepSpin::settings(int narg, char** arg) { aparam.push_back(atof(arg[iarg + 1 + ii])); } iarg += 1 + dim_aparam; + } else if (string(arg[iarg]) == string("charge_spin")) { + for (int ii = 0; ii < dim_chg_spin; ++ii) { + if (iarg + 1 + ii >= narg || is_key(arg[iarg + 1 + ii])) { + char tmp[1024]; + sprintf(tmp, "Illegal charge_spin, the dimension should be %d", + dim_chg_spin); + error->all(FLERR, tmp); + } + charge_spin.push_back(atof(arg[iarg + 1 + ii])); + } + iarg += 1 + dim_chg_spin; } else if (string(arg[iarg]) == string("ttm")) { #ifdef USE_TTM for (int ii = 0; ii < 1; ++ii) { diff --git a/source/lmp/tests/run_mpi_pair_deepmd_spin_graph_dpa4_pt2.py b/source/lmp/tests/run_mpi_pair_deepmd_spin_graph_dpa4_pt2.py new file mode 100644 index 0000000000..2eb01aca14 --- /dev/null +++ b/source/lmp/tests/run_mpi_pair_deepmd_spin_graph_dpa4_pt2.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Multi-rank LAMMPS driver for the native-spin DPA4 graph ``.pt2`` fixture. + +``atom_style spin`` / ``pair_style deepspin`` runner for +``test_lammps_dpa4_spin_graph_pt2.py``'s multi-rank comparison. The +native-spin DPA4 graph archive now carries the nested with-comm AOTI +artifact, so ``DeepSpinPTExpt::compute_inner`` drives a real domain- +decomposed run through ``run_model_graph_with_comm`` (per-block ghost +FEATURE refresh via ``border_op``; ghost SPINS arrive through the LAMMPS +``sp`` forward-comm). + +Rank 0 writes potential energy + per-atom force (3 cols) + per-atom +force_mag (3 cols) to ``OUTPUT``, id-ordered, so the parent pytest process +can compare a 2-rank run against a 1-rank run on the SAME archive. Same +output convention as ``run_mpi_pair_deepmd_spin_dpa3_pt2.py`` minus the +virial columns (the native-spin fixture takes no fparam/aparam). +""" + +from __future__ import ( + annotations, +) + +import argparse + +import numpy as np +from lammps import ( + PyLammps, +) +from mpi4py import ( + MPI, +) + +rank = MPI.COMM_WORLD.Get_rank() + +parser = argparse.ArgumentParser() +parser.add_argument( + "DATAFILE", type=str, help="LAMMPS data file (atom positions + spin)" +) +parser.add_argument("PB_FILE", type=str, help=".pt2 model file (native-spin graph)") +parser.add_argument("OUTPUT", type=str, help="Unused; kept for CLI-shape parity") +parser.add_argument( + "--processors", + type=str, + default="2 1 1", + help="LAMMPS processors grid. Default '2 1 1' forces multi-rank " + "domain decomposition (nswap>0). Pass '1 1 1' for a single-rank " + "reference run on the same archive.", +) +args = parser.parse_args() + +lammps = PyLammps() +lammps.processors(args.processors) +lammps.units("metal") +lammps.boundary("p p p") +lammps.atom_style("spin") +lammps.atom_modify("map yes") +lammps.neighbor("2.0 bin") +lammps.neigh_modify("every 10 delay 0 check no") +lammps.read_data(args.DATAFILE) +lammps.mass("1 58") +lammps.mass("2 16") +lammps.timestep(0.0005) +lammps.fix("1 all nve") +lammps.pair_style(f"deepspin {args.PB_FILE}") +lammps.pair_coeff("* *") +# Per-atom magnetic force components: LAMMPS does not expose ``fm`` through +# the legacy extract/gather_atoms registry, so go via +# ``compute property/atom fmx fmy fmz`` + ``gather``. +lammps.compute("fmprop all property/atom fmx fmy fmz") +lammps.run(0) + +forces_global = lammps.lmp.gather_atoms("f", 1, 3) +ids_global = lammps.lmp.gather_atoms("id", 0, 1) +fm_global = lammps.lmp.gather("c_fmprop", 1, 3) + +if rank == 0: + pe_global = lammps.eval("pe") + natoms = lammps.atoms.natoms + forces = np.array(forces_global, dtype=np.float64).reshape(natoms, 3) + fm = np.array(fm_global, dtype=np.float64).reshape(natoms, 3) + ids = np.array(ids_global, dtype=np.int64).reshape(natoms) + order = np.argsort(ids) + forces = forces[order] + fm = fm[order] + with open(args.OUTPUT, "w") as f: + f.write(f"{pe_global:.16e}\n") + # Each row: 3 force + 3 force_mag = 6 columns. + for fi, fmi in zip(forces, fm, strict=True): + row = np.concatenate([fi, fmi]) + f.write(" ".join(f"{v:.16e}" for v in row) + "\n") + +# Tear down LAMMPS before MPI.Finalize() so its destructor's MPI calls run +# while the communicator is still valid (see the dpa3 spin runner). +del lammps +MPI.Finalize() diff --git a/source/lmp/tests/test_lammps_dpa4_graph_pt2.py b/source/lmp/tests/test_lammps_dpa4_graph_pt2.py new file mode 100644 index 0000000000..4e7bd30980 --- /dev/null +++ b/source/lmp/tests/test_lammps_dpa4_graph_pt2.py @@ -0,0 +1,487 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Test LAMMPS with the NeighborGraph (graph-schema) .pt2 DPA4 model. + +The model ``deeppot_dpa4_graph.pt2`` is a DPA4/SeZM descriptor exported with +``lower_kind="graph"`` (gen_dpa4.py Section B). DPA4 is graph-native +end-to-end (no dense-only sub-block to gate off, unlike DPA2's +``use_three_body``), so the same config used for the dense ``deeppot_dpa4.pt2`` +fixture is graph-eligible; the weights are freshly jittered (see gen_dpa4.py +Section B.1) so the fixture is geometry-sensitive rather than the +architecturally edge-independent output of a fresh, untrained DPA4. + +DPA4's SeZM descriptor reads ghost-neighbour features at every interaction +block (``has_message_passing_across_ranks()`` returns ``self.bridging_switch +is None`` -- true for the plain (non-bridging) config exercised here), so +the GRAPH export auto-embeds a nested with-comm AOTI artifact +(``forward_lower_with_comm.pt2``) alongside the plain graph forward, the +same as DPA2's repformer. The DENSE (nlist) ``.pt2`` remains a single +comm-less artifact -- see ``test_lammps_dpa4_pt2.py``. + +Single-rank LAMMPS folds ghosts onto local owners (``fold_to_local=True``, +``N == nloc``) and uses the plain graph artifact, exactly as before. +Multi-rank LAMMPS keeps the extended region (``N == nall_real``) and routes +to the with-comm artifact: the C++ ``DeepPotPTExpt`` drives +``deepmd_export::border_op`` once per interaction block to exchange ghost +node/edge features across ranks, masks the fitting reduction to owned nodes +only, and LAMMPS reverse-comm folds the returned per-extended-atom forces +back onto their owners -- the same generic dispatch (``has_message_passing_`` ++ ``has_comm_artifact_`` + ``lower_input_kind == "graph"``) that already +serves DPA2, with zero C++ changes required for DPA4 (see +``source/api_cc/src/DeepPotPTExpt.cc``). + +Reference values come from ``source/tests/infer/gen_dpa4.py`` (the same +``deeppot_dpa4_graph.expected`` the C++ gtest uses). A second, independent +oracle -- ``deeppot_dpa4_graph_nlist_ref.pt2`` (same weights, dense-nlist +lower, NOT graph) -- is also exercised directly through LAMMPS so a +regression that only breaks the *C++* graph ingestion (not the Python +export path already cross-checked at gen-time) still gets caught. +""" + +import importlib.util +import os +import shutil +import signal +import subprocess as sp +import sys +import tempfile +from pathlib import ( + Path, +) + +import constants +import numpy as np +import pytest +from expected_ref import ( + read_expected_ref, +) +from lammps import ( + PyLammps, +) +from write_lmp_data import ( + write_lmp_data, +) + +pb_file = ( + Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot_dpa4_graph.pt2" +) +# Independent dense-nlist oracle exported from the SAME (jittered) weights +# (gen_dpa4.py B.1/B.2, lower_kind="nlist"); at non-binding sel graph and +# dense math are equivalent (gen-time cross-check already enforces atomic +# energy / force / total-virial agreement within 1e-8 at the Python level -- +# see gen_dpa4.py B.4). Comparing through LAMMPS as well exercises the C++ +# graph ingestion path (edge tensors, node atype slicing) independently of +# that Python-level check. +nlist_ref_file = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa4_graph_nlist_ref.pt2" +) +ref_file = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa4_graph.expected" +) +# The MPI runner is backend-agnostic (DATAFILE PB_FILE OUTPUT + flags); reuse +# the DPA3 driver verbatim rather than duplicate it (same pattern as +# test_lammps_dpa1_graph_pt2.py / test_lammps_dpa2_graph_pt2.py). +mpi_runner = Path(__file__).parent / "run_mpi_pair_deepmd_dpa3_pt2.py" + +# Ceiling for EVERY mpirun invocation (parse mode included): a with-comm +# desync hangs the collective forever, so an unbounded should-succeed +# regression would hang the whole suite. +_MPI_DEFAULT_TIMEOUT = 600.0 + +# Reference values written by source/tests/infer/gen_dpa4.py (PBC case). +# Guarded with try/except because gen_dpa4.py only runs when PyTorch is built, +# and the graph section is itself skipped under LeakSanitizer (see +# gen_dpa4.py Section B's module comment) -- either way this file must still +# be collectible, with the affected tests skipping cleanly. +try: + _ref = read_expected_ref(ref_file)["pbc"] + expected_e = float(np.sum(_ref["expected_e"])) + expected_f = _ref["expected_f"].reshape(6, 3) + # LAMMPS uses the opposite sign convention for virial vs DeepPot. + expected_v = -_ref["expected_v"].reshape(6, 9) +except FileNotFoundError: + expected_e = expected_f = expected_v = None + +_HAS_REF = expected_e is not None + +# Same 6-atom water system as the dense DPA4 fixture +# (source/tests/infer/gen_dpa4.py / test_lammps_dpa4_pt2.py): type_map +# [O, H], box 13x13x13. +box = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0]) +coord = np.array( + [ + [12.83, 2.56, 2.18], + [12.09, 2.87, 2.74], + [0.25, 3.32, 1.68], + [3.36, 3.00, 1.81], + [3.51, 2.51, 2.60], + [4.27, 3.22, 1.56], + ] +) +# Model type_map is ["O", "H"]; gen_dpa4.py atype = [0, 1, 1, 0, 1, 1] -> +# LAMMPS types [1, 2, 2, 1, 2, 2] under identity ``pair_coeff * *``. +type_OH = np.array([1, 2, 2, 1, 2, 2]) + +data_file = Path(__file__).parent / "data_dpa4_graph_pt2.lmp" +# Wide-box, 3-way x-split variant for the genuinely-empty-rank MPI corner +# (``processors 3 1 1``), same construction as +# ``test_lammps_dpa2_graph_pt2.py``'s ``data_file_empty_rank`` fixture but +# with DPA4's ghost cutoff: rcut(4.0)+skin(2.0)=6.0 (vs dpa2's 8.0). Atoms +# stay in x in [0.25, 12.83] near the left edge of a [0, 90] box. With 3 +# even x-slabs of width 30, rank 0 owns [0, 30) (all atoms), rank 2 owns +# [60, 90) (empty of local atoms but picks up a periodic ghost of the +# x~0.25 atoms wrapped around the box's x=90/x=0 seam, since that distance +# ~0.25 is well within the ghost cutoff), and rank 1 (the MIDDLE slab, +# [30, 60)) borders neither the real atoms directly (nearest real atom at +# distance 30-12.83 ~= 17.17 > 6) nor the periodic seam -- so rank 1 is the +# genuinely empty rank (zero owned AND zero ghost atoms) this fixture is +# built to produce. +data_file_empty_rank = Path(__file__).parent / "data_dpa4_graph_pt2_empty_rank.lmp" + + +def setup_module() -> None: + if os.environ.get("ENABLE_PYTORCH", "1") != "1": + pytest.skip( + "Skip test because PyTorch support is not enabled.", + ) + write_lmp_data(box, coord, type_OH, data_file) + box_empty_rank = np.array([0, 90, 0, 13, 0, 13, 0, 0, 0]) + write_lmp_data(box_empty_rank, coord, type_OH, data_file_empty_rank) + + +def teardown_module() -> None: + for f in [data_file, data_file_empty_rank]: + if f.exists(): + os.remove(f) + + +def _lammps(data_file, units="metal", atom_map: str = "yes") -> PyLammps: + lammps = PyLammps() + lammps.units(units) + lammps.boundary("p p p") + lammps.atom_style("atomic") + if atom_map != "no": + lammps.atom_modify(f"map {atom_map}") + lammps.neighbor("2.0 bin") + lammps.neigh_modify("every 10 delay 0 check no") + lammps.read_data(data_file.resolve()) + lammps.mass("1 16") + lammps.mass("2 2") + lammps.timestep(0.0005) + lammps.fix("1 all nve") + return lammps + + +@pytest.fixture +def lammps(): + lmp = _lammps(data_file=data_file) + yield lmp + lmp.close() + + +@pytest.mark.skipif( + not _HAS_REF, reason="gen_dpa4.py graph .expected fixture not generated" +) +def test_pair_deepmd(lammps) -> None: + """Single-rank serial run (``atom_modify map yes``): the graph .pt2 + folds ghosts onto local owners (``fold_to_local=True``) and must match + the gen_dpa4.py reference for energy and per-atom force. + """ + lammps.pair_style(f"deepmd {pb_file.resolve()}") + lammps.pair_coeff("* *") + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f[lammps.atoms[ii].id - 1] + ) + lammps.run(1) + + +@pytest.mark.skipif( + not _HAS_REF, reason="gen_dpa4.py graph .expected fixture not generated" +) +def test_pair_deepmd_virial(lammps) -> None: + """Single-rank per-atom virial via ``centroid/stress/atom``.""" + lammps.pair_style(f"deepmd {pb_file.resolve()}") + lammps.pair_coeff("* *") + lammps.compute("virial all centroid/stress/atom NULL pair") + for ii in range(9): + jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") + lammps.dump( + "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) + ) + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f[lammps.atoms[ii].id - 1] + ) + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 + for ii in range(9): + assert np.array( + lammps.variables[f"virial{ii}"].value + ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii]) + + +@pytest.mark.skipif( + not nlist_ref_file.exists(), + reason="gen_dpa4.py deeppot_dpa4_graph_nlist_ref.pt2 fixture not generated", +) +def test_pair_deepmd_graph_matches_nlist_ref() -> None: + """Single-rank graph .pt2 vs the independent dense-nlist oracle + (``deeppot_dpa4_graph_nlist_ref.pt2``, same weights, ``lower_kind="nlist"``) + through LAMMPS, energy/force within 1e-6. + + gen_dpa4.py already cross-checks graph vs nlist at the Python + ``DeepPot.eval`` level (B.4, 1e-8) at generation time; this test instead + drives BOTH artifacts through the C++ ``DeepPotPTExpt`` / LAMMPS pair + style, so a regression confined to the C++ graph-ingestion seam (edge + tensor construction, node atype slicing) that the Python-level gen-time + check cannot see is still caught. The per-atom virial is deliberately + NOT compared here: the graph path assigns each edge's force/virial + contribution fully to the source atom (edge_force_virial full-to-src), a + different (equally valid) decomposition than the dense path's -- only + energy and force (and, at the Python level already, the *total* virial) + are convention-independent. + """ + lmp_graph = _lammps(data_file=data_file) + lmp_graph.pair_style(f"deepmd {pb_file.resolve()}") + lmp_graph.pair_coeff("* *") + lmp_graph.run(0) + e_graph = lmp_graph.eval("pe") + f_graph = np.array([lmp_graph.atoms[ii].force for ii in range(6)]) + id_graph = [lmp_graph.atoms[ii].id for ii in range(6)] + lmp_graph.close() + + lmp_nlist = _lammps(data_file=data_file) + lmp_nlist.pair_style(f"deepmd {nlist_ref_file.resolve()}") + lmp_nlist.pair_coeff("* *") + lmp_nlist.run(0) + e_nlist = lmp_nlist.eval("pe") + f_nlist = np.array([lmp_nlist.atoms[ii].force for ii in range(6)]) + id_nlist = [lmp_nlist.atoms[ii].id for ii in range(6)] + lmp_nlist.close() + + # Same data file, single rank -> identical atom-id ordering; assert this + # rather than silently re-sorting so an ordering change is loud. + assert id_graph == id_nlist + assert e_graph == pytest.approx(e_nlist, rel=0, abs=1e-6) + np.testing.assert_allclose(f_graph, f_nlist, atol=1e-6, rtol=0) + + +# --------------------------------------------------------------------------- +# Multi-rank tests (message-passing with-comm graph route). +# +# DPA4's SeZM descriptor participates in per-block ghost exchange, so +# multi-rank LAMMPS routes to the nested with-comm artifact instead of the +# plain graph artifact used above. These tests are the correctness gate for +# that machinery on DPA4, mirroring ``test_lammps_dpa2_graph_pt2.py``'s +# multi-rank section (the SAME generic C++ dispatch; DPA4 required zero C++ +# changes, see the module docstring). +# --------------------------------------------------------------------------- + + +def _run_mpi_subprocess( + extra_args: list[str] | None = None, + nprocs: int = 2, + data_path: Path | None = None, + processors: str | None = None, + runner_args: list[str] | None = None, + pb: Path | None = None, + capture: bool = False, + timeout: float | None = None, +) -> dict: + """Invoke the (backend-agnostic) DPA3 MPI runner under + ``mpirun -n `` against the dpa4 graph .pt2 and return + ``{"pe": float, "forces": (n, 3), "virials": (n, 9)}``. + + ``nprocs == 1`` forces ``--processors 1 1 1`` so the C++ side sees + ``nprocs == 1`` and routes to the plain (single-rank) graph artifact -- + a same-archive reference for the multi-rank comparison. ``pb`` + overrides the model archive (defaults to ``deeppot_dpa4_graph.pt2``). + + EVERY run is bounded: ``timeout`` (seconds, default + ``_MPI_DEFAULT_TIMEOUT``) covers the parse path too, so a deadlocked + collective in a should-succeed regression cannot hang the suite -- on + expiry the WHOLE mpirun process group is SIGKILLed (killing only mpirun + can leave orphaned ranks blocking in the collective and holding the + GPU). With ``capture=True``, return raw subprocess info + (``returncode``, ``stdout``, ``stderr``, ``timed_out``) instead of + parsed output -- used by the fail-fast tests; a timeout there returns + ``timed_out=True`` with ``returncode=None`` for the caller to assert + on. In parse mode a timeout raises ``RuntimeError`` and a nonzero exit + raises ``subprocess.CalledProcessError`` (matching the old + ``check_call`` behavior). + """ + if data_path is None: + data_path = data_file + if pb is None: + pb = pb_file + if timeout is None: + timeout = _MPI_DEFAULT_TIMEOUT + with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f: + out_path = f.name + try: + argv = [ + "mpirun", + "-n", + str(nprocs), + sys.executable, + str(mpi_runner), + str(data_path.resolve()), + str(pb.resolve()), + out_path, + ] + if processors is not None: + argv.extend(["--processors", processors]) + elif nprocs == 1: + argv.extend(["--processors", "1 1 1"]) + if extra_args: + argv.extend(extra_args) + if runner_args: + argv.extend(runner_args) + proc = sp.Popen( + argv, + stdout=sp.PIPE if capture else None, + stderr=sp.PIPE if capture else None, + text=True, + start_new_session=True, + ) + try: + stdout, stderr = proc.communicate(timeout=timeout) + except sp.TimeoutExpired: + # Kill the whole process group: killing only mpirun can + # leave the deadlocked ranks orphaned (still blocking in + # the collective and holding the GPU). + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + stdout, stderr = proc.communicate() + if capture: + return { + "returncode": None, + "stdout": stdout or "", + "stderr": stderr or "", + "timed_out": True, + } + raise RuntimeError( + f"mpirun timed out after {timeout}s (process group killed); " + "a should-succeed MPI regression is deadlocked." + ) from None + if capture: + return { + "returncode": proc.returncode, + "stdout": stdout, + "stderr": stderr, + "timed_out": False, + } + if proc.returncode != 0: + raise sp.CalledProcessError(proc.returncode, argv) + with open(out_path) as fh: + lines = fh.read().strip().splitlines() + pe = float(lines[0]) + rows = np.array( + [list(map(float, line.split())) for line in lines[1:]], + dtype=np.float64, + ) + forces = rows[:, :3] + virials = rows[:, 3:] + return {"pe": pe, "forces": forces, "virials": virials} + finally: + if os.path.exists(out_path): + os.remove(out_path) + + +@pytest.mark.skipif( + shutil.which("mpirun") is None, reason="MPI is not installed on this system" +) +@pytest.mark.skipif( + importlib.util.find_spec("mpi4py") is None, reason="mpi4py is not installed" +) +def test_pair_deepmd_mpi_dpa4_graph_matches_single_rank() -> None: + """Multi-rank (``-n 2``) with-comm graph route must equal single-rank + (``-n 1``, plain graph artifact) on the SAME archive and trajectory. + + THE gate on the DPA4 with-comm graph machinery: per-block + ``deepmd_export::border_op`` ghost exchange, the owned-node energy + mask (extended region includes ghost nodes that must not contribute + to the reduced energy), and the reverse-comm force fold back onto + owners. A wrong-but-finite divergence in any of the three would show + up here even though there is no hardcoded reference value. + """ + out_mpi = _run_mpi_subprocess(nprocs=2) + out_ref = _run_mpi_subprocess(nprocs=1) + assert out_mpi["pe"] == pytest.approx(out_ref["pe"], rel=1e-8, abs=1e-10) + np.testing.assert_allclose(out_mpi["forces"], out_ref["forces"], atol=1e-8, rtol=0) + # Same tolerance as test_lammps_dpa2_graph_pt2.py's twin: a relative + # component absorbs the tiny ordering-dependent floating-point + # divergence the with-comm route's per-layer ghost exchange plus atomic + # index_add can introduce on CUDA, without loosening the CPU-exact + # (bit-reproducible) case. + np.testing.assert_allclose( + out_mpi["virials"], out_ref["virials"], atol=1e-8, rtol=1e-8 + ) + + +@pytest.mark.skipif( + shutil.which("mpirun") is None, reason="MPI is not installed on this system" +) +@pytest.mark.skipif( + importlib.util.find_spec("mpi4py") is None, reason="mpi4py is not installed" +) +def test_pair_deepmd_mpi_dpa4_graph_empty_rank_does_not_silently_succeed() -> None: + """A genuinely empty rank (zero owned AND zero ghost atoms) under the + message-passing with-comm graph route must NOT silently produce + wrong-but-plausible numbers. + + DPA4's with-comm route needs every rank to participate in the per-block + MPI ghost exchange (``border_op``); a rank with zero nodes has nothing + to export in the traced graph (violates the exported + ``Dim("n_node_total", min=1)`` and would desync the collective ghost + exchange across ranks). The C++ side (``DeepPotPTExpt.cc``, the SAME + model-agnostic guard already exercised by + ``test_lammps_dpa2_graph_pt2.py``) throws a clear, actionable error on + the empty rank instead of running. + + The failure is COLLECTIVE and PROMPT: before entering the per-layer + ``border_op`` collectives, every rank participates in a communicator- + wide min-reduction of its node count (``deepmd_export:: + allreduce_min_int``), so the non-empty peers detect the empty rank and + throw the same error instead of blocking forever waiting for it. A + timeout is therefore a FAILURE of this test (it would mean the + preflight regressed back into the historical deadlock), and the + documented error message must appear on a nonzero exit. + + ``data_file_empty_rank`` (3-way x-split, ``processors 3 1 1``) was + verified (see the module-level comment above the fixture) to put the + MIDDLE rank in a genuinely empty state, using DPA4's own ghost cutoff + (rcut(4.0)+skin(2.0)=6.0, vs dpa2's 8.0). + """ + out = _run_mpi_subprocess( + nprocs=3, + data_path=data_file_empty_rank, + processors="3 1 1", + capture=True, + timeout=120, + ) + assert not out["timed_out"], ( + "Multi-rank graph run with an empty rank timed out instead of " + "failing promptly: the collective empty-rank preflight " + "(allreduce_min_int) must make every rank throw BEFORE the " + "per-layer border_op collectives." + ) + assert out["returncode"] != 0, ( + "Expected the multi-rank message-passing graph run to fail loudly " + "on a genuinely empty rank, but it exited 0.\n" + f"stdout:\n{out['stdout'][-2000:]}\nstderr:\n{out['stderr'][-2000:]}" + ) + combined = out["stdout"] + out["stderr"] + assert "zero owned+ghost atoms" in combined, ( + "Expected the documented fail-loud message ('zero owned+ghost " + f"atoms'), got:\n{combined[-2000:]}" + ) diff --git a/source/lmp/tests/test_lammps_dpa4_pt2.py b/source/lmp/tests/test_lammps_dpa4_pt2.py index bf4fcf2940..1e47b1a9ed 100644 --- a/source/lmp/tests/test_lammps_dpa4_pt2.py +++ b/source/lmp/tests/test_lammps_dpa4_pt2.py @@ -7,32 +7,37 @@ Scope / coverage ---------------- ``deeppot_dpa4.pt2`` (generated by source/tests/infer/gen_dpa4.py) is a -DUAL-artifact archive: ``has_comm_artifact=True`` and -``has_message_passing=True`` (GNN). It is therefore the analogue of -``deeppot_dpa3_mpi.pt2`` (use_loc_mapping=False) — gen_dpa4.py does NOT -produce a separate use_loc_mapping=True archive, so the dpa3 cells A/B -(which need the no-with-comm .pt2) have no DPA4 counterpart. - -Single-rank cells covered here (all on the with-comm archive): - -- ``test_pair_deepmd`` — atom_modify map yes. Dispatch picks the regular - path because nswap==0 (single-rank PBC has an empty CommBrick - sendlist); the regular artifact uses the correct mapping built from the - LAMMPS atom-map. pe/forces must match the DeepPot reference. Mirrors - dpa3 cell C (``test_pair_deepmd_with_comm``). -- ``test_pair_deepmd_no_atom_map_fails_fast`` — atom_modify map no. - Despite the with-comm artifact being available, single-rank PBC has - nswap==0 so border_op cannot fill ghost features and the GNN model has - no reliable mapping. Must fail fast with the actionable - ``atom_modify map yes`` message. Mirrors dpa3 cell D - (``test_pair_deepmd_with_comm_no_atom_map_fails_fast``). +SINGLE-artifact archive: ``has_comm_artifact=False`` and +``has_message_passing=True`` (GNN). The DENSE (nlist) lower has no +comm_dict implementation (``dense_lower_supports_comm() == False``), so +``deserialize_to_file`` never emits a nested with-comm artifact for this +dense archive; multi-rank inference on it fails fast in C++ instead (see +the "Deferred" note below). This is scoped to the dense lower only: the +GRAPH ``.pt2`` (``deeppot_dpa4_graph.pt2``, ``test_lammps_dpa4_graph_pt2.py``) +DOES carry a with-comm artifact (``has_message_passing_across_ranks()`` +returns ``self.bridging_switch is None``) and runs multi-rank via the same +generic dispatch DPA2 uses. + +Single-rank cells covered here (all on this single artifact): + +- ``test_pair_deepmd`` — atom_modify map yes. The GNN model resolves + ghost-to-local mapping from the LAMMPS atom-map. pe/forces must match + the DeepPot reference. Mirrors dpa3 cell C + (``test_pair_deepmd_with_comm``). +- ``test_pair_deepmd_no_atom_map_fails_fast`` — atom_modify map no. The + GNN model has no reliable ghost-to-local mapping without the atom-map. + Must fail fast with the actionable ``atom_modify map yes`` message. + Mirrors dpa3 cell D (``test_pair_deepmd_with_comm_no_atom_map_fails_fast``). - virial / type_map / real-units / si-units variants mirror the dpa3 single-rank set. -Deferred (NOT covered): live multi-rank parity. DPA4 multi-rank -inference is out of PR-3 scope and has no mpi runner script. The C++ -with-comm dispatch is exercised for DPA4 only at the single-rank level -here; multi-rank DPA4 is left to a follow-up. +Deferred (NOT covered): live multi-rank parity for this DENSE archive. +``deeppot_dpa4.pt2`` carries no with-comm artifact, so multi-rank LAMMPS +inference on it fails fast at the first force evaluation instead of +running (see the module-level note); there is no with-comm dispatch to +exercise here, and no mpi runner script. Multi-rank inference for DPA4 IS +covered, but only through the GRAPH archive -- see +``test_lammps_dpa4_graph_pt2.py``'s multi-rank section. Tolerances match test_lammps_dpa3_pt2.py exactly (pytest.approx defaults for pe/forces; per-atom virial compared with pytest.approx). @@ -59,7 +64,8 @@ write_lmp_data, ) -# Dual-artifact (with-comm) DPA4 .pt2 — the only archive gen_dpa4.py emits. +# Single-artifact DPA4 .pt2 (has_comm_artifact=False) — the only archive +# gen_dpa4.py emits. pb_file = Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot_dpa4.pt2" ref_file = ( Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot_dpa4.expected" @@ -164,9 +170,9 @@ def lammps_no_atom_map(): def test_pair_deepmd(lammps) -> None: - # Single-rank with-comm archive + atom_modify map yes. Dispatch picks - # the regular path because nswap==0; the regular artifact uses the - # correct mapping built from the LAMMPS atom-map. Mirrors dpa3 cell C. + # Single-rank, single-artifact archive + atom_modify map yes. The GNN + # model resolves ghost-to-local mapping from the LAMMPS atom-map. + # pe/forces must match the DeepPot reference. Mirrors dpa3 cell C. lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.run(0) @@ -179,11 +185,10 @@ def test_pair_deepmd(lammps) -> None: def test_pair_deepmd_no_atom_map_fails_fast(lammps_no_atom_map) -> None: - # Single-rank with-comm archive + atom_modify map no. Single-rank PBC - # has an empty CommBrick sendlist (nswap==0), so border_op cannot fill - # ghost features and the GNN model has no reliable mapping. Must fail - # fast with the single-rank ``atom_modify map yes`` message. Mirrors - # dpa3 cell D. + # Single-rank, single-artifact archive + atom_modify map no. The GNN + # model has no reliable ghost-to-local mapping without the atom-map. + # Must fail fast with the single-rank ``atom_modify map yes`` message. + # Mirrors dpa3 cell D. lammps_no_atom_map.pair_style(f"deepmd {pb_file.resolve()}") lammps_no_atom_map.pair_coeff("* *") with pytest.raises(Exception, match=r"atom_modify map yes"): diff --git a/source/lmp/tests/test_lammps_dpa4_spin_graph_pt2.py b/source/lmp/tests/test_lammps_dpa4_spin_graph_pt2.py new file mode 100644 index 0000000000..96081f5f0d --- /dev/null +++ b/source/lmp/tests/test_lammps_dpa4_spin_graph_pt2.py @@ -0,0 +1,503 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Single-rank LAMMPS ``pair_style deepspin`` on the native-spin DPA4 +NeighborGraph (graph-schema) ``.pt2`` (Task 8's ``deeppot_dpa4_spin_graph.pt2``). + +Unlike the virtual-atom ``spin_ener`` scheme exercised by +``test_lammps_spin_pt2.py``, native-spin DPA4 has NO dense/nlist lower at +all -- spin rides the NeighborGraph lower exclusively (see +``deepmd/pt_expt/model/dpa4_native_spin_model.py``'s module docstring and +``source/tests/infer/gen_dpa4_spin.py``). The fixture also carries the +nested with-comm AOTI artifact, so multi-rank LAMMPS drives a real +domain-decomposed run through ``DeepSpinPTExpt::run_model_graph_with_comm``: +per-block ghost FEATURE refresh via ``border_op``, ghost SPINS via the +LAMMPS ``sp`` forward-comm (see ``source/api_cc/src/DeepSpinPTExpt.cc``). + +Reference (energy / force / force_mag / virial) values are computed LIVE at +test-setup time via ``deepmd.infer.DeepPot.eval`` on +``deeppot_dpa4_spin_graph.pt2`` for the fixed 4-atom NiO system reused from +``test_lammps_spin_pt2.py`` (box 13x13x13, same coordinates/type ordering: 2 +spin-active Ni + 2 non-magnetic O) -- i.e. exactly the Task 7 graph-spin +Python eval path, driven here through LAMMPS instead. + +The reference is deliberately NOT hardcoded (a previous revision hardcoded +it and went stale within ~1e-6 the moment master's ``dpa4_nn`` +physical-null-mass-attention change shifted DPA4 numerics -- exactly the +fragility flagged in the Task 10 review). Nor is it read from a sidecar +``.expected`` file produced by ``source/tests/infer/gen_dpa4_spin.py``: +that script's own PBC eval uses a DIFFERENT 6-atom (3 Ni + 3 O) system in a +6x6x6 box (see its module docstring / ``_COORDS`` / ``_CELL`` / ``_SPINS``), +and that box's edge length (6.0) exactly equals DPA4's ghost cutoff +(rcut(4.0)+skin(2.0)=6.0) -- not a safe geometry to reuse for a LAMMPS +periodic run. Instead, ``_compute_expected`` below loads the archive and +evaluates it, at test-setup time, on THIS module's own fixed geometry -- +mirroring ``test_lammps_model_devi_pt2.py``'s ``_compute_expected`` pattern +(subprocess-isolated, so importing ``deepmd``'s Python package does not +share a process with the LAMMPS plugin's own loaded ``libdeepmd_op_pt.so``). +This keeps the reference correct-by-construction: it always reflects +whatever the current archive produces, so a real DPA4 numerics shift is +caught by comparing against the *previous* run's output changing (reviewed +in the PR), not by a silently-stale hardcoded array. +""" + +import importlib.util +import json +import os +import shutil +import signal +import subprocess as sp +import sys +import tempfile +import textwrap +from pathlib import ( + Path, +) + +import constants +import numpy as np +import pytest +from lammps import ( + PyLammps, +) +from write_lmp_data import ( + write_lmp_data_spin, +) + +pb_file = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa4_spin_graph.pt2" +) +data_file = Path(__file__).parent / "data_dpa4_spin_graph_pt2.lmp" +# The MPI runner is graph-spin-specific (no aparam / no NULL-type +# extras, unlike run_mpi_pair_deepmd_spin_dpa3_pt2.py's virtual-atom-scheme +# runner): the native-spin DPA4 fixture takes no fparam/aparam. +mpi_runner = Path(__file__).parent / "run_mpi_pair_deepmd_spin_graph_dpa4_pt2.py" + +_MPI_DEFAULT_TIMEOUT = 120.0 + +# Same 4-atom NiO system as test_lammps_spin_pt2.py (box, coordinates, and +# LAMMPS type ordering all reused verbatim): 2 Ni atoms (LAMMPS type 1, +# deepmd atype 0, spin-active) + 2 O atoms (LAMMPS type 2, deepmd atype 1, +# non-magnetic) -- matches ``deeppot_dpa4_spin_graph.pt2``'s +# ``type_map=["Ni", "O"]`` and ``use_spin=[True, False]`` (gen_dpa4_spin.py). +box = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0]) +coord = np.array( + [ + [12.83, 2.56, 2.18], + [12.09, 2.87, 2.74], + [3.51, 2.51, 2.60], + [4.27, 3.22, 1.56], + ] +) +spin = np.array( + [ + [0, 0, 1.2737], + [0, 0, 1.2737], + [0, 0, 0], + [0, 0, 0], + ] +) +type_NiO = np.array([1, 1, 2, 2]) + +# LAMMPS's ``fm`` (what ``compute property/atom fmx fmy fmz`` reports) is +# NOT the raw DeepEval force_mag: pair_deepspin.cpp scales it by +# ``spin_norm / hbar`` per atom (metal-units ``hbar = 6.5821191e-04``, see +# ``source/lmp/pair_deepspin.cpp:531,535`` -- same convention already +# implicit, if untested against a raw LAMMPS ``fm`` read, in +# test_lammps_spin_pt2.py). ``spin_norm`` is 0 for the two non-magnetic O +# atoms, so the scaling is a no-op there (0 stays 0). +_HBAR_METAL = 6.5821191e-04 + +# Reference values (energy / atom-energy / force / force_mag / virial), +# populated by ``_compute_expected`` in ``setup_module`` -- see the module +# docstring for why these are computed live via a DeepPot subprocess call +# rather than hardcoded or read from a sidecar file. +expected_e = None +expected_ae = None +expected_f = None +expected_fm = None +expected_v = None + + +def _cell_from_lammps_box(lmp_box: np.ndarray) -> np.ndarray: + """Convert a LAMMPS ``xlo xhi ylo yhi zlo zhi xy xz yz`` box spec to a + flat, row-major 3x3 cell matrix (deepmd's ``box`` convention). + """ + xlo, xhi, ylo, yhi, zlo, zhi, xy, xz, yz = lmp_box + return np.array( + [ + xhi - xlo, + 0.0, + 0.0, + xy, + yhi - ylo, + 0.0, + xz, + yz, + zhi - zlo, + ] + ) + + +def _compute_expected() -> None: + """Load ``deeppot_dpa4_spin_graph.pt2`` via ``DeepPot`` and evaluate the + module's fixed 4-atom NiO system to obtain the Python reference. + + Runs in a subprocess to avoid importing ``deepmd`` in the LAMMPS test + process (see ``test_lammps_model_devi_pt2.py``'s ``_compute_expected`` + for the same precaution: the LAMMPS plugin already loads + ``libdeepmd_op_pt.so`` at the C++ level, and importing the Python + package on top of that can segfault). + """ + global expected_e, expected_ae, expected_f, expected_fm, expected_v + + cell = _cell_from_lammps_box(box) + atype = (type_NiO - 1).tolist() # LAMMPS 1-based -> deepmd 0-based (Ni=0, O=1) + + # ``deeppot_dpa4_spin_graph.pt2`` lives in ``source/tests/infer`` next to + # ``gen_common.py``, whose ``load_custom_ops()`` loads the build-tree + # ``libdeepmd_op_pt.so`` (registering ``deepmd::edge_force_virial``, which + # the graph ``.pt2`` inference needs). ``import deepmd.pt`` alone only loads + # the op library from SHARED_LIB_DIR, which the build-test env does not + # populate -- so the subprocess reuses that fallback (after importing + # ``deepmd.pt``, per its docstring) before constructing ``DeepPot``. + infer_dir = str(pb_file.resolve().parent) + script = textwrap.dedent(f"""\ + import json + import sys + import numpy as np + + sys.path.insert(0, {infer_dir!r}) + import deepmd.pt # noqa: F401 (triggers the base op-library load) + from gen_common import load_custom_ops + + load_custom_ops() + from deepmd.infer import DeepPot + + dp = DeepPot({str(pb_file.resolve())!r}) + e, f, v, ae, av, fm, mm = dp.eval( + np.array({coord.tolist()!r}).reshape(1, -1, 3), + np.array({cell.tolist()!r}).reshape(1, 9), + {atype!r}, + atomic=True, + spin=np.array({spin.tolist()!r}).reshape(1, -1, 3), + ) + print(json.dumps({{ + "e": float(e[0, 0]), + "ae": np.asarray(ae[0]).reshape(-1).tolist(), + "f": np.asarray(f[0]).tolist(), + "fm": np.asarray(fm[0]).tolist(), + "av": np.asarray(av[0]).tolist(), + }})) + """) + proc = sp.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + ) + if proc.returncode != 0: + raise RuntimeError(f"Failed to compute expected values:\n{proc.stderr}") + result = json.loads(proc.stdout.strip()) + + expected_e = result["e"] + expected_ae = np.array(result["ae"]) + expected_f = np.array(result["f"]) + # Raw DeepEval force_mag (dE/dspin), scaled by LAMMPS's own + # spin_norm / hbar unit convention (see the comment on ``_HBAR_METAL`` + # above) before comparison. + fm_raw = np.array(result["fm"]) + spin_norm = np.linalg.norm(spin, axis=1) + expected_fm = fm_raw * (spin_norm / _HBAR_METAL)[:, None] + # Per-atom virial, sign-flipped (LAMMPS convention) relative to DeepPot's + # atomic virial output (mirrors test_lammps_spin_pt2.py's convention). + expected_v = -np.array(result["av"]) + + +def setup_module() -> None: + if os.environ.get("ENABLE_PYTORCH", "1") != "1": + pytest.skip( + "Skip test because PyTorch support is not enabled.", + ) + if not pb_file.exists(): + pytest.skip("deeppot_dpa4_spin_graph.pt2 not found") + _compute_expected() + write_lmp_data_spin(box, coord, spin, type_NiO, data_file) + + +def teardown_module() -> None: + if data_file.exists(): + os.remove(data_file) + + +def _lammps(data_file, units="metal") -> PyLammps: + """Standard DeepSpin LAMMPS system, plus ``atom_modify map yes``. + + Mirrors ``lammps_test_utils.make_spin_lammps`` (not reused directly: it + does not set ``atom_modify``), with the map turned on -- the native-spin + DPA4 GRAPH ``.pt2`` needs the LAMMPS atom-map to resolve ghost-atom + indices to local owners for single-rank inference (same requirement as + the energy graph route; see ``pair_deepspin.cpp``'s + ``DeePMD-kit Error: Single-rank LAMMPS .pt2 inference requires + `atom_modify map yes``` check). + """ + if units != "metal": + raise ValueError("units for spin should be metal") + + lammps = PyLammps() + lammps.units(units) + lammps.boundary("p p p") + lammps.atom_style("spin") + lammps.atom_modify("map yes") + lammps.neighbor("2.0 bin") + lammps.neigh_modify("every 10 delay 0 check no") + lammps.read_data(data_file.resolve()) + lammps.mass("1 58") + lammps.mass("2 16") + lammps.timestep(0.0005) + lammps.fix("1 all nve") + return lammps + + +@pytest.fixture +def lammps(): + lmp = _lammps(data_file=data_file) + yield lmp + lmp.close() + + +def _gather_force_mag(lammps: PyLammps, natoms: int) -> np.ndarray: + """Extract per-atom force_mag in atom-id order. + + LAMMPS does not expose ``fm`` through the legacy ``extract``/ + ``gather_atoms`` registry (see ``run_mpi_pair_deepmd_spin_dpa3_pt2.py``'s + module docstring), so go via ``compute property/atom fmx fmy fmz`` + + ``gather`` (id-ordered on every rank, single-rank included). + """ + fm_global = lammps.lmp.gather("c_fmprop", 1, 3) + return np.array(fm_global, dtype=np.float64).reshape(natoms, 3) + + +def test_pair_deepspin(lammps) -> None: + """Single-rank LAMMPS energy + force + force_mag vs the Python DeepEval + graph-spin reference (Task 7 path), on the same 4-atom NiO system. + """ + lammps.pair_style(f"deepspin {pb_file.resolve()}") + lammps.pair_coeff("* *") + lammps.compute("fmprop all property/atom fmx fmy fmz") + lammps.run(0) + + assert lammps.eval("pe") == pytest.approx(expected_e) + + forces = np.array([lammps.atoms[ii].force for ii in range(4)], dtype=np.float64) + ids = np.array([lammps.atoms[ii].id for ii in range(4)]) + order = np.argsort(ids) + forces = forces[order] + np.testing.assert_allclose(forces, expected_f, atol=1e-8, rtol=0) + + force_mag = _gather_force_mag(lammps, coord.shape[0]) + np.testing.assert_allclose(force_mag, expected_fm, atol=1e-8, rtol=0) + # Anti-vacuity / native-spin design invariant: force_mag on the two + # non-spin (O) atoms must be exactly zero, both in the Python reference + # (baked into expected_fm above) and as produced by LAMMPS. + np.testing.assert_array_equal(force_mag[2:], np.zeros((2, 3))) + + lammps.run(1) + + +def test_pair_deepspin_virial(lammps) -> None: + """Single-rank per-atom pe/pressure/virial via + ``pe/atom`` / ``pressure`` / ``centroid/stress/atom``, atol=1e-8, + rtol=1e-8. + """ + lammps.pair_style(f"deepspin {pb_file.resolve()}") + lammps.pair_coeff("* *") + lammps.compute("peatom all pe/atom pair") + lammps.compute("pressure all pressure NULL pair") + lammps.compute("virial all centroid/stress/atom NULL pair") + lammps.variable("eatom atom c_peatom") + for ii in range(9): + jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] + lammps.variable(f"pressure{jj} equal c_pressure[{ii + 1}]") + for ii in range(9): + jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") + lammps.dump( + "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) + ) + lammps.run(0) + + assert lammps.eval("pe") == pytest.approx(expected_e) + + forces = np.array([lammps.atoms[ii].force for ii in range(4)], dtype=np.float64) + ids = np.array([lammps.atoms[ii].id for ii in range(4)]) + order = np.argsort(ids) + forces = forces[order] + np.testing.assert_allclose(forces, expected_f, atol=1e-8, rtol=0) + + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 + np.testing.assert_allclose( + np.array(lammps.variables["eatom"].value), + expected_ae[idx_map], + atol=1e-8, + rtol=1e-8, + ) + + vol = box[1] * box[3] * box[5] + for ii in range(6): + jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] + pressure_jj = np.array(lammps.variables[f"pressure{jj}"].value) / ( + constants.nktv2p + ) + expected_pressure_jj = -expected_v[idx_map, jj].sum(axis=0) / vol + np.testing.assert_allclose( + pressure_jj, expected_pressure_jj, atol=1e-8, rtol=1e-8 + ) + for ii in range(9): + jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] + virial_jj = np.array(lammps.variables[f"virial{jj}"].value) / (constants.nktv2p) + np.testing.assert_allclose( + virial_jj, expected_v[idx_map, jj], atol=1e-8, rtol=1e-8 + ) + + +# --------------------------------------------------------------------------- +# Multi-rank: the native-spin graph .pt2 carries the nested with-comm +# artifact, so a 2-rank run must REPRODUCE the 1-rank result (energy, +# force and force_mag) rather than fail fast. +# --------------------------------------------------------------------------- + + +def _run_mpi_subprocess( + extra_args: list[str] | None = None, + nprocs: int = 2, + data_path: Path | None = None, + processors: str | None = None, + capture: bool = False, + timeout: float | None = None, +) -> dict: + """Invoke the graph-spin MPI runner under ``mpirun -n `` against + the native-spin DPA4 graph ``.pt2``. + + Copied (module-global closure, not imported) from + ``test_lammps_dpa4_graph_pt2.py``'s twin. With ``capture=True``, return + raw subprocess info (``returncode``, ``stdout``, ``stderr``, + ``timed_out``) -- used by the fail-fast test below; every invocation is + bounded by ``timeout`` (default ``_MPI_DEFAULT_TIMEOUT``) so a + should-fail-but-doesn't run cannot hang the suite, and on expiry the + WHOLE mpirun process group is SIGKILLed. + """ + if data_path is None: + data_path = data_file + if timeout is None: + timeout = _MPI_DEFAULT_TIMEOUT + with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f: + out_path = f.name + try: + argv = [ + "mpirun", + "-n", + str(nprocs), + sys.executable, + str(mpi_runner), + str(data_path.resolve()), + str(pb_file.resolve()), + out_path, + ] + if processors is not None: + argv.extend(["--processors", processors]) + elif nprocs == 1: + argv.extend(["--processors", "1 1 1"]) + if extra_args: + argv.extend(extra_args) + proc = sp.Popen( + argv, + stdout=sp.PIPE if capture else None, + stderr=sp.PIPE if capture else None, + text=True, + start_new_session=True, + ) + try: + stdout, stderr = proc.communicate(timeout=timeout) + except sp.TimeoutExpired: + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + stdout, stderr = proc.communicate() + if capture: + return { + "returncode": None, + "stdout": stdout or "", + "stderr": stderr or "", + "timed_out": True, + } + raise RuntimeError( + f"mpirun timed out after {timeout}s (process group killed); " + "a should-succeed MPI regression is deadlocked." + ) from None + if capture: + return { + "returncode": proc.returncode, + "stdout": stdout, + "stderr": stderr, + "timed_out": False, + } + if proc.returncode != 0: + raise sp.CalledProcessError(proc.returncode, argv) + with open(out_path) as fh: + lines = fh.read().strip().splitlines() + pe = float(lines[0]) + rows = np.array( + [list(map(float, line.split())) for line in lines[1:]], + dtype=np.float64, + ) + return {"pe": pe, "rows": rows} + finally: + if os.path.exists(out_path): + os.remove(out_path) + + +@pytest.mark.skipif( + shutil.which("mpirun") is None, reason="MPI is not installed on this system" +) +@pytest.mark.skipif( + importlib.util.find_spec("mpi4py") is None, reason="mpi4py is not installed" +) +def test_pair_deepspin_mpi_matches_single_rank() -> None: + """A 2-rank MPI run must reproduce the 1-rank result on the SAME archive. + + The native-spin DPA4 graph ``.pt2`` now carries the nested with-comm + artifact, so ``DeepSpinPTExpt::compute_inner`` drives a real + domain-decomposed run through ``run_model_graph_with_comm``: the + per-block ghost FEATURE refresh rides ``border_op`` while ghost SPINS + arrive via the LAMMPS ``sp`` forward-comm. Both the conservative force + and the MAGNETIC force must be rank-count invariant -- force_mag is the + output that only exists on this route, so comparing it is what proves + the spin leaf survived the with-comm lower. + + Replaces the previous fail-fast test, which asserted the C++ throw that + existed only while native spin was excluded from the with-comm export. + """ + single = _run_mpi_subprocess(nprocs=1, processors="1 1 1") + multi = _run_mpi_subprocess(nprocs=2, processors="2 1 1") + + # anti-vacuity: a degenerate fixture (all-zero forces) would make the + # comparison pass for the wrong reason. + assert np.abs(multi["rows"][:, :3]).max() > 1e-6, "forces are trivially zero" + assert np.abs(multi["rows"][:, 3:6]).max() > 1e-6, "force_mag is trivially zero" + + np.testing.assert_allclose( + multi["pe"], single["pe"], rtol=1e-10, atol=1e-10, err_msg="energy" + ) + np.testing.assert_allclose( + multi["rows"][:, :3], + single["rows"][:, :3], + rtol=1e-10, + atol=1e-10, + err_msg="force", + ) + np.testing.assert_allclose( + multi["rows"][:, 3:6], + single["rows"][:, 3:6], + rtol=1e-10, + atol=1e-10, + err_msg="force_mag", + ) diff --git a/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py b/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py new file mode 100644 index 0000000000..135bf04426 --- /dev/null +++ b/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py @@ -0,0 +1,391 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Single-rank LAMMPS ``pair_style deepmd`` on the ZBL-BRIDGED DPA4 +NeighborGraph (graph-schema) ``.pt2`` (``deeppot_dpa4_zbl_graph.pt2``, +``source/tests/infer/gen_dpa4_zbl.py``). + +``bridging_method: ZBL`` builds a COMPOSITION -- ``LinearEnergyModel`` over +``[learned DPA4, InterPotentialAtomicModel]`` with ``weights="sum"`` -- so +this drives the graph lower of a linear composition through the LAMMPS pair +style. The archive already had a C++ gtest +(``source/api_cc/tests/test_deeppot_dpa4_zbl_ptexpt.cc``) but NO LAMMPS +coverage at all; the pair style is a distinct consumer (its own nlist/ghost +handling, per-atom virial accumulation and unit conversion), so a regression +confined to it was invisible. + +Single-rank only, deliberately +------------------------------ +Bridging enables the descriptor's Source Freeze Propagation Gate, whose +per-node ``eta_j = prod_{e: src_e = j} w_e`` folds a node's FULL outgoing-edge +set. Edges exist only for owned centres, so eta is incomplete on every rank +and the freeze exports NO with-comm artifact (``gen_dpa4_zbl.py`` asserts +``has_comm_artifact is False`` and that no nested +``forward_lower_with_comm.pt2`` entry exists). There is therefore no +correct multi-rank answer to compare against; what this file pins instead is +that a multi-rank run FAILS LOUDLY rather than silently returning +wrong-but-plausible numbers -- see +``test_pair_deepmd_mpi_dpa4_zbl_fails_fast``. + +Reference values are computed LIVE at test-setup time via +``deepmd.infer.DeepPot.eval`` on the archive itself, mirroring +``test_lammps_dpa4_spin_graph_pt2.py``'s ``_compute_expected`` (which explains +the reasoning in full). Two reasons, both load-bearing here: + +- A hardcoded array goes stale the moment DPA4 numerics shift, and this + fixture's energies are dominated by a ~700 eV analytical term, so a stale + reference would fail in a way that looks like a bridging bug. +- ``gen_dpa4_zbl.py``'s own ``.expected`` sidecar cannot be reused: its + evaluation uses a 6x6x6 A cell whose edge length equals DPA4's LAMMPS ghost + cutoff exactly (rcut(4.0) + skin(2.0) = 6.0), which is not a safe geometry + for a periodic LAMMPS run. This module keeps the same 6-atom geometry (the + 0.9 A Ni-Ni pair that makes the ZBL term dominant) in a 13x13x13 A box + instead. +""" + +import importlib.util +import json +import os +import shutil +import signal +import subprocess as sp +import sys +import tempfile +from pathlib import ( + Path, +) + +import constants +import numpy as np +import pytest +from lammps import ( + PyLammps, +) +from write_lmp_data import ( + write_lmp_data, +) + +pb_file = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa4_zbl_graph.pt2" +) +data_file = Path(__file__).parent / "data_dpa4_zbl_pt2.lmp" +# The MPI runner is backend-agnostic (DATAFILE PB_FILE OUTPUT + flags); reuse +# the DPA3 driver verbatim rather than duplicate it (same pattern as +# test_lammps_dpa4_graph_pt2.py). Only the fail-fast test below uses it. +mpi_runner = Path(__file__).parent / "run_mpi_pair_deepmd_dpa3_pt2.py" + +# Ceiling for the mpirun invocation. The fail-fast under test is expected to +# throw on EVERY rank before any collective, so a timeout means the guard +# regressed into a deadlock -- which is a test failure, not a slow machine. +_MPI_DEFAULT_TIMEOUT = 300.0 + +# 6-atom NiO system, coordinates verbatim from +# ``source/tests/infer/gen_dpa4_zbl.py``'s ``_COORDS``: atoms 0 and 1 sit +# 0.9 A apart, inside ``bridging_r_outer``, so the analytical ZBL term +# contributes a large, unmistakable repulsion (~1.4e3 eV/A forces) rather +# than a numerical afterthought. The box is 13x13x13 A (NOT the generator's +# 6x6x6 -- see the module docstring). +box = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0]) +coord = np.array( + [ + [1.0, 1.0, 1.0], + [1.9, 1.0, 1.0], + [1.3, 1.8, 1.0], + [0.4, 1.2, 1.6], + [3.6, 2.0, 1.3], + [3.4, 0.7, 1.7], + ] +) +# Model ``type_map`` is ["Ni", "O"]; gen_dpa4_zbl.py's atype [0,0,0,1,1,1] +# -> LAMMPS types [1,1,1,2,2,2] under identity ``pair_coeff * *``. +type_NiO = np.array([1, 1, 1, 2, 2, 2]) + +# Reference values, populated by ``_compute_expected`` in ``setup_module``. +expected_e = None +expected_ae = None +expected_f = None +expected_v = None + + +def _cell_from_lammps_box(lmp_box: np.ndarray) -> np.ndarray: + """Convert a LAMMPS ``xlo xhi ylo yhi zlo zhi xy xz yz`` box spec to a + flat, row-major 3x3 cell matrix (deepmd's ``box`` convention). + """ + xlo, xhi, ylo, yhi, zlo, zhi, xy, xz, yz = lmp_box + return np.array([xhi - xlo, 0.0, 0.0, xy, yhi - ylo, 0.0, xz, yz, zhi - zlo]) + + +def _compute_expected() -> None: + """Load ``deeppot_dpa4_zbl_graph.pt2`` via ``DeepPot`` and evaluate this + module's fixed 6-atom system to obtain the Python reference. + + Runs in a subprocess to avoid importing ``deepmd`` in the LAMMPS test + process (the LAMMPS plugin already loads ``libdeepmd_op_pt.so`` at the C++ + level, and importing the Python package on top of that can segfault) -- + the same precaution as ``test_lammps_dpa4_spin_graph_pt2.py`` and + ``test_lammps_model_devi_pt2.py``. + """ + global expected_e, expected_ae, expected_f, expected_v + + cell = _cell_from_lammps_box(box) + atype = (type_NiO - 1).tolist() # LAMMPS 1-based -> deepmd 0-based + + # The archive lives in ``source/tests/infer`` next to ``gen_common.py``, + # whose ``load_custom_ops()`` loads the build-tree ``libdeepmd_op_pt.so`` + # (registering ``deepmd::edge_force_virial``, which graph ``.pt2`` + # inference needs). ``import deepmd.pt`` alone only loads the op library + # from SHARED_LIB_DIR, which the build-test env does not populate. + infer_dir = str(pb_file.resolve().parent) + script = ( + "import json, sys\n" + "import numpy as np\n" + f"sys.path.insert(0, {infer_dir!r})\n" + "import deepmd.pt # noqa: F401 (triggers the base op-library load)\n" + "from gen_common import load_custom_ops\n" + "load_custom_ops()\n" + "from deepmd.infer import DeepPot\n" + f"dp = DeepPot({str(pb_file.resolve())!r})\n" + "e, f, v, ae, av = dp.eval(\n" + f" np.array({coord.tolist()!r}).reshape(1, -1, 3),\n" + f" np.array({cell.tolist()!r}).reshape(1, 9),\n" + f" {atype!r},\n" + " atomic=True,\n" + ")\n" + "print(json.dumps({\n" + ' "e": float(e[0, 0]),\n' + ' "ae": np.asarray(ae[0]).reshape(-1).tolist(),\n' + ' "f": np.asarray(f[0]).tolist(),\n' + ' "av": np.asarray(av[0]).tolist(),\n' + "}))\n" + ) + proc = sp.run([sys.executable, "-c", script], capture_output=True, text=True) + if proc.returncode != 0: + raise RuntimeError(f"Failed to compute expected values:\n{proc.stderr}") + result = json.loads(proc.stdout.strip()) + + expected_e = result["e"] + expected_ae = np.array(result["ae"]) + expected_f = np.array(result["f"]) + # LAMMPS uses the opposite sign convention for the virial vs DeepPot. + expected_v = -np.array(result["av"]) + + # Anti-vacuity, checked once here so every test below is known to compare + # against a non-degenerate reference: the 0.9 A Ni-Ni pair must drive a + # large analytical ZBL repulsion. A fresh (unjittered) DPA4 would give + # identically zero forces, and a composition that lost its analytical + # child would give small ones. + assert np.max(np.abs(expected_f)) > 1e2, ( + "the ZBL-bridged reference forces are too small to be the analytical " + f"term on a 0.9 A pair (max |f| = {np.max(np.abs(expected_f)):.3e}); " + "the fixture or the bridging composition is degenerate." + ) + + +def setup_module() -> None: + if os.environ.get("ENABLE_PYTORCH", "1") != "1": + pytest.skip("Skip test because PyTorch support is not enabled.") + if not pb_file.exists(): + pytest.skip( + "deeppot_dpa4_zbl_graph.pt2 not found (run " + "source/tests/infer/gen_dpa4_zbl.py)." + ) + _compute_expected() + write_lmp_data(box, coord, type_NiO, data_file) + + +def teardown_module() -> None: + if data_file.exists(): + os.remove(data_file) + + +def _lammps(data_file, units="metal") -> PyLammps: + """Standard LAMMPS system plus ``atom_modify map yes``. + + DPA4 message-passes within a rank, so the single-rank graph ``.pt2`` + route needs the LAMMPS atom-map to resolve ghost-atom indices to their + local owners (``DeepPotPTExpt.cc``: "Single-rank LAMMPS .pt2 inference + requires `atom_modify map yes`"). + """ + lammps = PyLammps() + lammps.units(units) + lammps.boundary("p p p") + lammps.atom_style("atomic") + lammps.atom_modify("map yes") + lammps.neighbor("2.0 bin") + lammps.neigh_modify("every 10 delay 0 check no") + lammps.read_data(data_file.resolve()) + lammps.mass("1 58") # Ni + lammps.mass("2 16") # O + lammps.timestep(0.0005) + lammps.fix("1 all nve") + return lammps + + +@pytest.fixture +def lammps(): + lmp = _lammps(data_file=data_file) + yield lmp + lmp.close() + + +def test_pair_deepmd(lammps) -> None: + """Single-rank energy + per-atom force vs the Python DeepEval reference. + + ``rel=1e-10`` on the energy and ``atol=1e-8`` on the force: both sides are + fp64 and run the SAME compiled artifact, so this is a cross-consumer + (LAMMPS pair style vs Python DeepEval) check rather than a cross-backend + one. The absolute force bound is ~1e-11 relative at this fixture's ~1.4e3 + eV/A magnitude, matching the bound the sibling DPA4 graph LAMMPS tests + use. + """ + lammps.pair_style(f"deepmd {pb_file.resolve()}") + lammps.pair_coeff("* *") + lammps.run(0) + + assert lammps.eval("pe") == pytest.approx(expected_e, rel=1e-10) + + ids = np.array([lammps.atoms[ii].id for ii in range(6)]) + forces = np.array([lammps.atoms[ii].force for ii in range(6)]) + np.testing.assert_allclose(forces, expected_f[ids - 1], atol=1e-8, rtol=0) + + # A second MD step: the ZBL repulsion is large but bounded (~0.03 A of + # displacement at dt = 0.5 fs), so this is a stability smoke test of the + # per-step dispatch, not a dynamics check. + lammps.run(1) + + +def test_pair_deepmd_atom_energy_and_virial(lammps) -> None: + """Single-rank per-atom energy and per-atom virial. + + ``centroid/stress/atom`` is the pair style's own virial accumulation + path, which the C++ gtest does not exercise; the atomic energies pin that + the composition's two per-atom terms are summed per atom rather than only + in the total. + """ + lammps.pair_style(f"deepmd {pb_file.resolve()}") + lammps.pair_coeff("* *") + lammps.compute("peatom all pe/atom pair") + lammps.compute("virial all centroid/stress/atom NULL pair") + lammps.variable("eatom atom c_peatom") + for ii in range(9): + jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") + lammps.dump( + "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) + ) + lammps.run(0) + + assert lammps.eval("pe") == pytest.approx(expected_e, rel=1e-10) + + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 + np.testing.assert_allclose( + np.array(lammps.variables["eatom"].value), + expected_ae[idx_map], + atol=1e-8, + rtol=0, + ) + for ii in range(9): + np.testing.assert_allclose( + np.array(lammps.variables[f"virial{ii}"].value) / constants.nktv2p, + expected_v[idx_map, ii], + atol=1e-8, + rtol=0, + ) + + +# --------------------------------------------------------------------------- +# Multi-rank: NOT a correctness test -- a fail-fast test. +# --------------------------------------------------------------------------- + + +def _run_mpi_subprocess(nprocs: int, processors: str, timeout: float) -> dict: + """Run the (backend-agnostic) DPA3 MPI runner against the bridged archive + and return ``{"returncode", "stdout", "stderr", "timed_out"}``. + + Always bounded: on expiry the WHOLE mpirun process group is SIGKILLed + (killing only mpirun can leave orphaned ranks blocking in a collective). + """ + with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f: + out_path = f.name + try: + argv = [ + "mpirun", + "-n", + str(nprocs), + sys.executable, + str(mpi_runner), + str(data_file.resolve()), + str(pb_file.resolve()), + out_path, + "--processors", + processors, + ] + proc = sp.Popen( + argv, stdout=sp.PIPE, stderr=sp.PIPE, text=True, start_new_session=True + ) + try: + stdout, stderr = proc.communicate(timeout=timeout) + except sp.TimeoutExpired: + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + stdout, stderr = proc.communicate() + return { + "returncode": None, + "stdout": stdout or "", + "stderr": stderr or "", + "timed_out": True, + } + return { + "returncode": proc.returncode, + "stdout": stdout, + "stderr": stderr, + "timed_out": False, + } + finally: + if os.path.exists(out_path): + os.remove(out_path) + + +@pytest.mark.skipif( + shutil.which("mpirun") is None, reason="MPI is not installed on this system" +) +@pytest.mark.skipif( + importlib.util.find_spec("mpi4py") is None, reason="mpi4py is not installed" +) +def test_pair_deepmd_mpi_dpa4_zbl_fails_fast() -> None: + """A multi-rank run of a BRIDGED archive must fail loudly, not answer. + + The bridged model is single-rank only by construction (see the module + docstring), so its freeze exports no with-comm artifact while still + declaring ``has_message_passing``. ``DeepPotPTExpt::compute_inner``'s + dispatch reads exactly that combination -- graph lower + ``nprocs > 1`` + + message passing + no with-comm artifact -- and throws before building any + tensors. Without the guard the run would fall through to the plain + single-rank artifact on a per-rank subdomain, where the bridging gate's + per-node eta is incomplete: wrong, finite, plausible numbers. + + The failure is uniform across ranks (every rank evaluates the same + metadata-only predicate before any collective), so a TIMEOUT is a failure + of this test: it would mean the guard regressed into a deadlock. + + This is deliberately the ONLY multi-rank test in this file; there is no + correct multi-rank reference for a bridged model to compare against. + """ + out = _run_mpi_subprocess( + nprocs=2, processors="2 1 1", timeout=_MPI_DEFAULT_TIMEOUT + ) + assert not out["timed_out"], ( + "Multi-rank run of the bridged archive timed out instead of failing " + "promptly; the dispatch guard must throw on every rank BEFORE any " + "collective." + ) + assert out["returncode"] != 0, ( + "Expected the multi-rank run of a bridged (no with-comm artifact) " + "archive to fail loudly, but it exited 0.\n" + f"stdout:\n{out['stdout'][-2000:]}\nstderr:\n{out['stderr'][-2000:]}" + ) + combined = out["stdout"] + out["stderr"] + assert "with-comm artifact" in combined, ( + "Expected the documented fail-loud message (mentioning the missing " + f"'with-comm artifact'), got:\n{combined[-2000:]}" + ) diff --git a/source/tests/common/dpmodel/test_descrpt_dpa4.py b/source/tests/common/dpmodel/test_descrpt_dpa4.py index a687fffcb8..d2b238ae8e 100644 --- a/source/tests/common/dpmodel/test_descrpt_dpa4.py +++ b/source/tests/common/dpmodel/test_descrpt_dpa4.py @@ -94,12 +94,20 @@ def test_shapes_and_interface(self) -> None: def test_message_passing_semantics(self) -> None: # SeZM always resolves ghost neighbours on the lower path, so it always - # reports message passing; cross-rank ghost exchange is needed only when - # zone bridging is disabled (a BridgingSwitch cannot be reproduced by a - # single rank for ghost owners). + # reports message passing. The GRAPH lower implements the cross-rank + # exchange via a real per-layer border_op, so a plain (non-bridging) + # descriptor reports across_ranks True; its DENSE lower has no + # comm_dict implementation (the dense adapter raises on it), so + # dense_lower_supports_comm() is False and the freeze machinery + # skips the dead dense with-comm artifact. Source Freeze Propagation + # bridging is excluded from across_ranks: its per-node gate folds a + # node's entire outgoing-edge set, which a single rank cannot + # observe for ghost owners, so bridging models fail fast on + # multi-rank instead. dd = make_descriptor() assert dd.has_message_passing() is True assert dd.has_message_passing_across_ranks() is True + assert dd.dense_lower_supports_comm() is False dd_bridge = make_descriptor(inner_clamp_r_inner=0.5, inner_clamp_r_outer=1.0) assert dd_bridge.has_message_passing() is True assert dd_bridge.has_message_passing_across_ranks() is False @@ -115,6 +123,24 @@ def test_serialize_roundtrip_exact(self) -> None: out2 = np.asarray(dd2.call(coord.reshape(nf, -1), atype, nlist)[0]) np.testing.assert_array_equal(out1, out2) + def test_random_gamma_inference_deterministic(self) -> None: + """The dpmodel backend never applies the random local-Z roll, even when configured. + + ``random_gamma`` is a training-only augmentation gated by the + ``_in_training_mode`` runtime hook; dpmodel has no training mode, so + the hook is ``False`` and two calls of a ``random_gamma=True`` + descriptor are bit-identical (the pt_expt twin's train-mode + behavior is pinned in + ``source/tests/pt_expt/descriptor/test_dpa4.py::test_random_gamma_train_eval_gate``). + """ + dd = make_descriptor(random_gamma=True) + assert dd._in_training_mode() is False + coord, atype, nlist = make_inputs() + nf = atype.shape[0] + out1 = np.asarray(dd.call(coord.reshape(nf, -1), atype, nlist)[0]) + out2 = np.asarray(dd.call(coord.reshape(nf, -1), atype, nlist)[0]) + np.testing.assert_array_equal(out1, out2) + def test_permutation_equivariance(self) -> None: dd = make_descriptor() coord, atype, nlist = make_inputs() diff --git a/source/tests/common/dpmodel/test_dp_atomic_model.py b/source/tests/common/dpmodel/test_dp_atomic_model.py index 721a4f5895..468c2c6834 100644 --- a/source/tests/common/dpmodel/test_dp_atomic_model.py +++ b/source/tests/common/dpmodel/test_dp_atomic_model.py @@ -55,6 +55,11 @@ def test_methods(self) -> None: self.assertEqual(md0.get_dim_aparam(), 0) self.assertEqual(md0.mixed_types(), ds.mixed_types()) self.assertEqual(md0.get_sel_type(), [0, 1]) + # Base-default (False) branch of the descriptor spin capabilities, + # cached at construction via direct method calls (the True branch is + # pinned in test_dpa4_call_graph.py). + self.assertFalse(md0.supports_native_spin()) + self.assertFalse(md0.supports_charge_spin) def test_self_consistency( self, diff --git a/source/tests/common/dpmodel/test_dpa1_graph_model_energy.py b/source/tests/common/dpmodel/test_dpa1_graph_model_energy.py index e03dcaa840..da0649cfda 100644 --- a/source/tests/common/dpmodel/test_dpa1_graph_model_energy.py +++ b/source/tests/common/dpmodel/test_dpa1_graph_model_energy.py @@ -267,3 +267,12 @@ def test_graph_lower_invariant_to_charge_spin() -> None: assert with_cs[k] is None else: np.testing.assert_array_equal(with_cs[k], v) + + +def test_graph_type_embedding_table_matches_type_embedding() -> None: + # The seam hook must return exactly the descriptor's full tebd table. + dd = _make_model([200]).get_descriptor() + np.testing.assert_array_equal( + np.asarray(dd.graph_type_embedding_table()), + np.asarray(dd.type_embedding.call()), + ) diff --git a/source/tests/common/dpmodel/test_dpa4_call_graph.py b/source/tests/common/dpmodel/test_dpa4_call_graph.py new file mode 100644 index 0000000000..9896049594 --- /dev/null +++ b/source/tests/common/dpmodel/test_dpa4_call_graph.py @@ -0,0 +1,977 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Graph-route regression tests for the dpmodel DPA4 descriptor.""" + +import dataclasses + +import numpy as np +import pytest + +from deepmd.dpmodel.atomic_model import ( + DPAtomicModel, +) +from deepmd.dpmodel.descriptor.dpa4 import ( + DescrptDPA4, + _graph_from_padded_nlist, +) +from deepmd.dpmodel.fitting.dpa4_ener import ( + SeZMEnergyFittingNet, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, +) +from deepmd.dpmodel.utils.nlist import ( + extend_input_and_build_neighbor_list, +) + +from ...dpa4_fixtures import ( + jitter_zero_arrays, +) + + +def build_neighbor_list_np(coord, rcut, nnei): + """Build a padded, distance-sorted gas-phase neighbor list (no PBC). + + Parameters + ---------- + coord + Coordinates with shape (nf, nloc, 3). + rcut + Cutoff radius. + nnei + Number of neighbor slots; pads with -1. + + Returns + ------- + np.ndarray + Neighbor list with shape (nf, nloc, nnei) holding local indices. + """ + nf, nloc, _ = coord.shape + nlist = -np.ones((nf, nloc, nnei), dtype=np.int64) + for f in range(nf): + dist = np.linalg.norm(coord[f][:, None, :] - coord[f][None, :, :], axis=-1) + for i in range(nloc): + neighbors = [ + (dist[i, j], j) for j in range(nloc) if j != i and dist[i, j] < rcut + ] + neighbors.sort() + for slot, (_, j) in enumerate(neighbors[:nnei]): + nlist[f, i, slot] = j + return nlist + + +def build_sparse_edges_from_nlist(coord, nlist): + """Extract the valid physical edges of a padded neighbor list. + + The padded layout keeps one slot per neighbor (``-1`` marks padding). The + graph-route edge contract -- edges packed into a :class:`NeighborGraph` + and consumed by ``call_graph`` -- is one explicit edge per kept slot, + indexing the flattened frame-major node axis (``node = f * nloc + i``). + The edge vector points from the center toward the neighbor, matching the + padded path's ``r_j - r_i``. + + Parameters + ---------- + coord + Coordinates with shape (nf, nloc, 3). + nlist + Neighbor list with shape (nf, nloc, nnei); -1 marks padding. + + Returns + ------- + tuple[np.ndarray, np.ndarray] + ``edge_index`` with shape (2, E) (rows are src, dst) and ``edge_vec`` + with shape (E, 3), aligned on the same edge axis in row-major + ``(frame, center, slot)`` order. + """ + nf, nloc, nnei = nlist.shape + src, dst, vec = [], [], [] + for f in range(nf): + for i in range(nloc): + for s in range(nnei): + j = int(nlist[f, i, s]) + if j < 0: + continue + src.append(f * nloc + j) + dst.append(f * nloc + i) + vec.append(coord[f, j] - coord[f, i]) + edge_index = np.asarray([src, dst], dtype=np.int64) # (2, E) + edge_vec = np.asarray(vec, dtype=np.float64) # (E, 3) + return edge_index, edge_vec + + +def make_descriptor() -> DescrptDPA4: + return DescrptDPA4( + ntypes=3, + sel=8, + rcut=4.0, + channels=16, + n_radial=8, + lmax=2, + mmax=1, + n_blocks=2, + precision="float64", + seed=7, + random_gamma=False, + ) + + +def make_inputs(seed=7, nf=2, nloc=6, rcut=4.0, nnei=8, ntypes=3): + rng = np.random.default_rng(seed) + coord = rng.uniform(0.0, 3.5, size=(nf, nloc, 3)) + atype = rng.integers(0, ntypes, size=(nf, nloc)) + nlist = build_neighbor_list_np(coord, rcut, nnei) + return coord, atype, nlist + + +def make_graph_from_nlist(coord, nlist): + """Build a ghost-free NeighborGraph from a gas-phase local nlist.""" + nf, nloc, _ = nlist.shape + edge_index, edge_vec = build_sparse_edges_from_nlist(coord, nlist) + return NeighborGraph( + n_node=np.full((nf,), nloc, dtype=np.int64), + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=np.ones(edge_index.shape[1], dtype=bool), + ) + + +def _run_graph(dd, coord, atype, nlist, permute_seed=None): + graph = make_graph_from_nlist(coord, nlist) + if permute_seed is not None: + perm = np.random.default_rng(permute_seed).permutation( + graph.edge_index.shape[1] + ) + graph = dataclasses.replace( + graph, + edge_index=graph.edge_index[:, perm], + edge_vec=graph.edge_vec[perm], + edge_mask=graph.edge_mask[perm], + ) + out, rot_mat = dd.call_graph(graph, atype.reshape(-1)) + assert rot_mat is None + return np.asarray(out) + + +def make_message_sensitive_descriptor(seed: int = 99) -> DescrptDPA4: + """A ``make_descriptor()`` variant with its zero-init residuals jittered. + + Two calls with the same ``seed`` are bit-identical (deserialize is + deterministic given the jittered parameter tree), so a pair of + independently constructed descriptors used to isolate the effect of + ``exclude_types`` still share every other weight. + """ + data = make_descriptor().serialize() + data = jitter_zero_arrays(data, np.random.default_rng(seed)) + return DescrptDPA4.deserialize(data) + + +def make_spin_descriptor(seed: int = 99) -> DescrptDPA4: + """A ``use_spin`` variant of ``make_message_sensitive_descriptor()``. + + Native spin (``use_spin`` on type 0) is enabled on top of the same + zero-init-residual jitter used by ``make_message_sensitive_descriptor``, + so the descriptor is sensitive to both edges/messages AND spin -- a bare + ``make_descriptor()`` variant would be architecturally spin-independent + for the same reason it is edge-independent (see that helper's + docstring), which would make a spin-sensitivity check vacuous. + """ + dd = DescrptDPA4( + ntypes=3, + sel=8, + rcut=4.0, + channels=16, + n_radial=8, + lmax=2, + mmax=1, + n_blocks=2, + precision="float64", + seed=7, + random_gamma=False, + use_spin=[True, False, False], + ) + data = dd.serialize() + data = jitter_zero_arrays(data, np.random.default_rng(seed)) + return DescrptDPA4.deserialize(data) + + +def test_call_graph_spin_sensitivity() -> None: + """call_graph(spin=...) must change the output (teeth: spin reaches the trunk).""" + dd = make_spin_descriptor() + coord, atype, nlist = make_inputs() + graph = make_graph_from_nlist(coord, nlist) + atype_flat = atype.reshape(-1) + rng = np.random.default_rng(3) + spin = rng.normal(size=(atype_flat.shape[0], 3)) + out0, _ = dd.call_graph(graph, atype_flat) + out1, _ = dd.call_graph(graph, atype_flat, spin=spin) + assert not np.allclose(np.asarray(out0), np.asarray(out1)) + + +def test_call_graph_spin_matches_dense_adapter() -> None: + """Graph-lower spin path == dense adapter spin path (shared trunk, 1e-12).""" + dd = make_spin_descriptor() + coord, atype, nlist = make_inputs() + nf, nloc = atype.shape + rng = np.random.default_rng(3) + spin = rng.normal(size=(nf * nloc, 3)) + ref, *_ = dd.call( + coord.reshape(nf, -1), atype, nlist, spin=spin.reshape(nf, nloc, 3) + ) + graph, atype_flat = _graph_from_padded_nlist(coord, atype, nlist, None) + out, _ = dd.call_graph(graph, atype_flat, spin=spin) + np.testing.assert_allclose( + np.asarray(out).reshape(nf, nloc, -1), + np.asarray(ref), + rtol=1e-12, + atol=1e-12, + ) + + +def make_charge_spin_descriptor(seed: int = 99) -> DescrptDPA4: + """An ``add_chg_spin_ebd`` variant of ``make_message_sensitive_descriptor()``. + + Charge/spin FiLM conditions the type embedding directly (see + ``_apply_charge_spin_embedding``), which -- per + ``jitter_zero_arrays``'s docstring -- IS the fresh descriptor's scalar + read-out (zero-init residuals make the blocks near-identity), so even a + bare ``make_descriptor(add_chg_spin_ebd=True)`` would already be + charge_spin-sensitive. Jittered anyway for consistency with every other + graph-route fixture in this module (project convention: all + sensitivity/parity fixtures jitter). + """ + dd = DescrptDPA4( + ntypes=3, + sel=8, + rcut=4.0, + channels=16, + n_radial=8, + lmax=2, + mmax=1, + n_blocks=2, + precision="float64", + seed=7, + random_gamma=False, + add_chg_spin_ebd=True, + ) + data = dd.serialize() + data = jitter_zero_arrays(data, np.random.default_rng(seed)) + return DescrptDPA4.deserialize(data) + + +def test_call_graph_charge_spin_sensitivity() -> None: + """call_graph(charge_spin=...) must change the output (teeth: charge_spin reaches the trunk).""" + dd = make_charge_spin_descriptor() + coord, atype, nlist = make_inputs() + graph = make_graph_from_nlist(coord, nlist) + atype_flat = atype.reshape(-1) + nf = atype.shape[0] + cs0 = np.tile(np.array([[0.0, 1.0]]), (nf, 1)) + cs1 = np.tile(np.array([[1.0, 1.0]]), (nf, 1)) + out0, _ = dd.call_graph(graph, atype_flat, charge_spin=cs0) + out1, _ = dd.call_graph(graph, atype_flat, charge_spin=cs1) + assert not np.allclose(np.asarray(out0), np.asarray(out1)) + + +def test_call_graph_charge_spin_matches_dense_adapter() -> None: + """Graph-lower charge_spin path == dense adapter charge_spin path (shared trunk, 1e-12). + + Uses DISTINCT charge_spin values per frame (nf=2) so the test also pins + the per-frame (nf, 2) -> flat-N association: a wrong nf threaded into + ``_apply_charge_spin_embedding`` would either shape-error or silently mix + the two frames' conditioning. + """ + dd = make_charge_spin_descriptor() + coord, atype, nlist = make_inputs() + nf, nloc = atype.shape + cs = np.array([[0.3, -0.7], [1.0, 0.2]]) + ref, *_ = dd.call(coord.reshape(nf, -1), atype, nlist, charge_spin=cs) + graph, atype_flat = _graph_from_padded_nlist(coord, atype, nlist, None) + out, _ = dd.call_graph(graph, atype_flat, charge_spin=cs) + np.testing.assert_allclose( + np.asarray(out).reshape(nf, nloc, -1), + np.asarray(ref), + rtol=1e-12, + atol=1e-12, + ) + + +def test_call_graph_matches_dense() -> None: + # Same physical edges (non-binding sel) => same descriptor within fp64 + # scatter-reassociation tolerance; output is flat (N, C). + # + # NOTE: uses make_message_sensitive_descriptor(), not the plain + # make_descriptor() fixture -- see that helper's docstring. A bare + # make_descriptor() is architecturally edge-independent (multiple + # zero-init residual output projections), so this parity check would + # pass trivially (0.0 == 0.0) regardless of whether call_graph's edge + # handling is correct. + dd = make_message_sensitive_descriptor() + coord, atype, nlist = make_inputs() + nf, nloc = atype.shape + out_dense = np.asarray(dd.call(coord.reshape(nf, -1), atype, nlist)[0]) + out_graph = _run_graph(dd, coord, atype, nlist) + assert out_graph.shape == (nf * nloc, dd.get_dim_out()) + np.testing.assert_allclose( + out_graph.reshape(nf, nloc, -1), out_dense, rtol=1e-10, atol=1e-12 + ) + + +def test_call_graph_matches_dense_permuted_edges() -> None: + # Arbitrary edge order (the graph contract) must not change the result. + # + # NOTE: uses make_message_sensitive_descriptor() -- see + # test_call_graph_matches_dense's NOTE above; a bare make_descriptor() + # is edge-independent, so permuting its (irrelevant) edges is vacuous. + dd = make_message_sensitive_descriptor() + coord, atype, nlist = make_inputs() + nf, nloc = atype.shape + out_dense = np.asarray(dd.call(coord.reshape(nf, -1), atype, nlist)[0]) + out_graph = _run_graph(dd, coord, atype, nlist, permute_seed=31) + np.testing.assert_allclose( + out_graph.reshape(nf, nloc, -1), out_dense, rtol=1e-10, atol=1e-12 + ) + + +def test_message_sensitive_fixture_is_edge_dependent() -> None: + """Pin that make_message_sensitive_descriptor() is edge-dependent. + + The two parity tests above (and the exclude_types test below) are only + meaningful because make_message_sensitive_descriptor() jitters DPA4's + zero-init residual projections so the output actually depends on + edges/messages -- a bare make_descriptor() does not (see + jitter_zero_arrays's docstring in dpa4_fixtures.py). This test durably + pins that edge-sensitivity: it runs call_graph once as-is and once with every + edge masked out, and asserts the outputs differ by a non-trivial + margin. If a future change to the fixture (or to DPA4's zero-init + scheme) silently made it edge-independent again, this test fails loud + instead of the parity tests above going quietly vacuous. + """ + dd = make_message_sensitive_descriptor() + coord, atype, nlist = make_inputs() + graph = make_graph_from_nlist(coord, nlist) + graph_no_edges = dataclasses.replace( + graph, edge_mask=np.zeros_like(graph.edge_mask) + ) + out_with_edges, _ = dd.call_graph(graph, atype.reshape(-1)) + out_no_edges, _ = dd.call_graph(graph_no_edges, atype.reshape(-1)) + out_with_edges = np.asarray(out_with_edges) + out_no_edges = np.asarray(out_no_edges) + assert np.max(np.abs(out_with_edges - out_no_edges)) > 1e-6 + + +def test_call_graph_exclude_types_matches_dense() -> None: + # Pair exclusion: canonical apply_pair_exclusion on the graph must equal + # the dense build_type_exclude_mask route. Also pins the empty-exclusion + # branch via the tests above. + # + # NOTE: uses make_message_sensitive_descriptor(), not the plain + # make_descriptor() fixture, so the anti-vacuity assertion below is + # meaningful -- see that helper's docstring. A bare make_descriptor() + # is architecturally edge-independent (multiple zero-init residual + # output projections), so exclude_types provably cannot change its + # output; asserting non-vacuity against it would always fail regardless + # of whether exclusion is correctly wired. + dd = make_message_sensitive_descriptor() + dd.reinit_exclude([(0, 1)]) + coord, atype, nlist = make_inputs() + nf, nloc = atype.shape + out_dense = np.asarray(dd.call(coord.reshape(nf, -1), atype, nlist)[0]) + out_graph = _run_graph(dd, coord, atype, nlist) + np.testing.assert_allclose( + out_graph.reshape(nf, nloc, -1), out_dense, rtol=1e-10, atol=1e-12 + ) + # anti-vacuity: exclusion must actually change the descriptor + dd2 = make_message_sensitive_descriptor() + out_noexcl = np.asarray(dd2.call(coord.reshape(nf, -1), atype, nlist)[0]) + assert not np.allclose(out_dense, out_noexcl, rtol=1e-6, atol=1e-8) + + +def test_call_graph_comm_dict_raises() -> None: + dd = make_descriptor() + coord, atype, nlist = make_inputs() + graph = make_graph_from_nlist(coord, nlist) + with pytest.raises(NotImplementedError, match="comm_dict"): + dd.call_graph(graph, atype.reshape(-1), comm_dict={"dummy": None}) + + +def test_capability_flags() -> None: + dd = make_descriptor() + assert dd.uses_graph_lower() is True + assert dd.uses_compact_edge_pairs() is False + assert dd.graph_type_embedding_table() is None + dd.disable_graph_lower() + assert dd.uses_graph_lower() is False + + +def test_supports_native_spin_capability_gate() -> None: + """Pin DPA4's ``True`` branch of the spin-capability contract. + + ``supports_native_spin``/``supports_charge_spin`` are declared on + ``BaseDescriptor`` (``make_base_descriptor``) with a concrete default of + ``False``; DPA4 overrides both to ``True`` and ``DPAtomicModel`` caches + the answers via direct method calls in ``__init__``. The + inherited-default (``False``) branch is pinned in + ``test_base_descriptor_capabilities.py``. + """ + dd = make_descriptor() + assert dd.supports_native_spin() is True + assert dd.supports_charge_spin() is True + ft = SeZMEnergyFittingNet( + ntypes=3, + dim_descrpt=dd.get_dim_out(), + neuron=[16], + precision="float64", + seed=5, + ) + dpa4_model = DPAtomicModel(dd, ft, type_map=["A", "B", "C"]) + assert dpa4_model.supports_native_spin() is True + assert dpa4_model.supports_charge_spin is True + + +def test_uses_graph_lower_feature_gates() -> None: + # Every conditioning input DPA4 supports now rides the graph lower: native + # spin, charge/spin FiLM, and SFPG bridging. Only the explicit escape + # hatch (disable_graph_lower / _graph_lower_disabled) gates it off. + for attr in ("charge_spin_embedding", "bridging_switch", "spin_embedding"): + dd = make_descriptor() + assert dd.uses_graph_lower() is True + setattr(dd, attr, object()) # any non-None sentinel + assert dd.uses_graph_lower() is True, attr + dd = make_descriptor() + assert dd.uses_graph_lower() is True + dd.disable_graph_lower() + assert dd.uses_graph_lower() is False + + +def test_dpa4_ener_fitting_call_graph_matches_dense() -> None: + # The inherited flat-N call_graph must be bit-identical to the dense + # call for the custom GLU fitting nets. + from deepmd.dpmodel.fitting.dpa4_ener import ( + SeZMEnergyFittingNet, + ) + + rng = np.random.default_rng(11) + ntypes, nf, nloc, nd = 3, 2, 6, 16 + ft = SeZMEnergyFittingNet( + ntypes=ntypes, + dim_descrpt=nd, + neuron=[24, 24], + precision="float64", + seed=5, + ) + dd = rng.standard_normal((nf, nloc, nd)) + atype = rng.integers(0, ntypes, size=(nf, nloc)) + ref = np.asarray(ft(dd, atype)["energy"]) + got = np.asarray( + ft.call_graph(dd.reshape(nf * nloc, nd), atype.reshape(-1))["energy"] + ) + np.testing.assert_allclose(got.reshape(ref.shape), ref, rtol=1e-12, atol=1e-14) + + +def make_bridging_descriptor(seed: int = 99) -> DescrptDPA4: + """A SFPG-bridging variant of ``make_message_sensitive_descriptor()``. + + ``inner_clamp_r_inner``/``inner_clamp_r_outer`` build ``InnerClamp`` (edge + distance freeze) and ``BridgingSwitch`` (per-source edge gate) -- see + ``test_message_passing_semantics`` in ``test_descrpt_dpa4.py`` for the + same construction. Jittered for the same message-sensitivity reason as + ``make_message_sensitive_descriptor``. + """ + dd = DescrptDPA4( + ntypes=2, + sel=8, + rcut=4.0, + channels=16, + n_radial=8, + lmax=2, + mmax=1, + n_blocks=2, + precision="float64", + seed=7, + random_gamma=False, + inner_clamp_r_inner=0.8, + inner_clamp_r_outer=1.2, + ) + data = dd.serialize() + data = jitter_zero_arrays(data, np.random.default_rng(seed)) + return DescrptDPA4.deserialize(data) + + +def _sphere_points(n_points: int, radius: float) -> np.ndarray: + """Deterministic Fibonacci-like sphere sampling around the origin.""" + idx = np.arange(n_points, dtype=np.float64) + phi = np.pi * (3.0 - np.sqrt(5.0)) # golden-angle step + z = 1.0 - 2.0 * (idx + 0.5) / float(n_points) + rho = np.sqrt(np.clip(1.0 - z * z, 0.0, None)) + theta = phi * idx + return np.stack( + [radius * rho * np.cos(theta), radius * rho * np.sin(theta), radius * z], + axis=-1, + ) # (n_points, 3) + + +def _frozen_sphere_descriptor_outputs( + dd: DescrptDPA4, near_distance: float, n_points: int = 12 +) -> np.ndarray: + """3-atom probe (A, B, C): A fixed at the origin, B rigidly slides on a + sphere of radius ``near_distance`` around A (inside the frozen zone), C + anchored well outside the bridging window as an ordinary GNN neighbor of + both. Returns the flat call_graph output reshaped to (n_points, 3, C). + + Mirrors pt's ``TestSourceFreezePropagationGate._build_three_atom_box`` / + ``_evaluate_frozen_sphere_atom_energies``, at the descriptor level (no + fitting net needed: if the descriptor feature of A is invariant, any + downstream per-atom fitting net -- itself a function of A's fixed + descriptor and fixed type only -- is trivially invariant too). + """ + directions = _sphere_points(n_points, near_distance) + coord = np.zeros((n_points, 3, 3), dtype=np.float64) + coord[:, 1, :] = directions # B rotates around A, radius fixed + coord[:, 2, :] = np.array([2.4, 0.0, 0.0]) # C: ordinary neighbor, static + atype = np.tile(np.array([[0, 1, 0]], dtype=np.int64), (n_points, 1)) + nlist = build_neighbor_list_np(coord, dd.get_rcut(), nnei=4) + graph = make_graph_from_nlist(coord, nlist) + out, _ = dd.call_graph(graph, atype.reshape(-1)) + return np.asarray(out).reshape(n_points, 3, -1) + + +def test_call_graph_bridging_frozen_sphere_invariance() -> None: + """SFPG bridging on the graph route: A's descriptor must be invariant to + the rigid motion of its frozen partner B (inside r_inner=0.8). + """ + dd = make_bridging_descriptor() + out = _frozen_sphere_descriptor_outputs(dd, near_distance=0.5) + span_a = np.max(np.abs(out[:, 0, :] - out[0:1, 0, :])) + assert span_a < 1e-10, span_a + + +def test_call_graph_bridging_leak_reopens_when_gate_disabled() -> None: + """Ablation: clearing bridging_switch (InnerClamp stays active) must + reopen the direction/multi-hop leak on the graph route -- pins that SFPG, + not InnerClamp alone, owns the invariance above. + """ + dd = make_bridging_descriptor() + dd.bridging_switch = None + out = _frozen_sphere_descriptor_outputs(dd, near_distance=0.5) + span_a = np.max(np.abs(out[:, 0, :] - out[0:1, 0, :])) + assert span_a > 1e-6, span_a + + +def test_charge_spin_model_routes_through_graph_lower() -> None: + """A native charge_spin DPA4 model must reach ``call_graph`` (not the old + ``cs -> dense`` gate) when a graph ``neighbor_graph_method`` is requested, + and the output must still be charge_spin-sensitive. + """ + from unittest.mock import ( + patch, + ) + + from deepmd.dpmodel.model.model import ( + get_model, + ) + + config = { + "type_map": ["Ni", "O"], + "descriptor": { + "type": "dpa4", + "rcut": 4.0, + "sel": 8, + "channels": 16, + "n_radial": 8, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "precision": "float64", + "seed": 7, + "random_gamma": False, + "add_chg_spin_ebd": True, + }, + "fitting_net": {"type": "dpa4_ener", "neuron": [8, 8]}, + } + model = get_model(config) + data = model.serialize() + data = jitter_zero_arrays(data, np.random.default_rng(17)) + model = type(model).deserialize(data) + + rng = np.random.default_rng(23) + nf, nloc = 1, 6 + coord = rng.uniform(0.5, 3.5, size=(nf, nloc, 3)) + atype = rng.integers(0, 2, size=(nf, nloc)) + box = 8.0 * np.eye(3, dtype=np.float64)[None] + cs0 = np.array([[0.0, 1.0]]) + cs1 = np.array([[1.0, 1.0]]) + + # Spy on the INSTANCE's bound method (wrapping the bound method), not the + # class with autospec: ``autospec=True`` + ``wraps=`` mis-binds + # ``self`` across Python/mock versions (green on 3.13, "not enough values + # to unpack" on 3.10 CI). Patching the instance and wrapping its bound + # method is version-stable. + descriptor = model.atomic_model.descriptor + with patch.object( + descriptor, + "call_graph", + wraps=descriptor.call_graph, + ) as spy: + out0 = model.call_common( + coord, atype, box, charge_spin=cs0, neighbor_graph_method="dense" + ) + assert spy.call_count >= 1, ( + "charge_spin must not force the model back onto the dense lower" + ) + out1 = model.call_common( + coord, atype, box, charge_spin=cs1, neighbor_graph_method="dense" + ) + assert not np.allclose(out0["energy_redu"], out1["energy_redu"]) + + +# Golden values pinned at the pre-reroute dense ``call`` (Task 7 controller +# Step 1.5). Generated at commit 12fe36ce (before the dense body became a +# NeighborGraph adapter) on an edge-sensitive descriptor. These MUST stay green +# after the reroute: a failure is a real adapter bug (edge-enumeration order, +# masking semantics, or precision-cast placement) -- never regenerate them and +# never loosen the tolerance. +_GOLDEN_CALL_DENSE_A = np.array( + [ + -0.08188957220398312, + -1.640109878466157, + 0.8980097395843406, + 0.6061720338823796, + 0.0369348802764473, + -0.4605616073317207, + 0.2663720748033271, + -1.0019825204289972, + 0.9408144687458523, + -0.6099639089284248, + 2.4055685430467664, + 0.6359644283661099, + 0.6903588394569702, + -0.9224964930776121, + 0.4990734769035464, + -1.0289550441210629, + -0.052560049562915594, + -1.17182802800252, + -0.4399327261943537, + 0.9966067341202732, + 1.5687440250412155, + -3.0306808891197763, + 1.3242653552348236, + -0.9019250803105836, + 1.252395311571487, + -0.514322347326694, + 3.040128995033731, + 0.45798849935105596, + -1.3940405513593772, + 0.8850637923045088, + 1.0372907246516128, + -0.4515792038039226, + -1.6828699160972829, + -0.01753501653910832, + 1.3975042415458003, + 0.8305783891840123, + 0.017013409258949445, + -0.047541417689215784, + 1.7317295338479273, + 0.2646210000936154, + 1.7083324210702961, + -0.458134646915977, + 2.59754428069986, + -0.12463014845722649, + -0.6443660208642956, + -2.071356246344078, + 2.5475184391276517, + -3.261424696743016, + 0.7055742966569417, + 0.5939618410270757, + -0.56485826918787, + 1.0694437511308956, + 0.37928349246064613, + -1.5092229530421246, + 1.693119227855766, + -0.11732801749609194, + -0.3825884526212911, + -0.4160475055486212, + 0.6259876087807682, + 0.7919342525125217, + -0.927223820144914, + -0.5152672801477078, + -0.23177697201360192, + -0.6026480324401201, + -0.7302717757588707, + -2.0898379534230243, + -0.24268458779685936, + 1.0001761022724518, + 1.6824564427885655, + -2.6335349596487907, + 0.8333379970412093, + -1.10986823930664, + 1.96965257621882, + -1.2096194291700268, + 3.8541594037592026, + 0.6124194446383112, + -0.9192570952710638, + 0.6938046842943415, + 1.846103645926, + -1.0138859880807036, + -0.28339562628465337, + -1.217166056202735, + 1.1321978569539521, + 0.3540199393504897, + -0.010795152497377874, + -0.43097091905968254, + -0.17835626273011676, + -0.5311505123026412, + 1.3411564031588632, + -0.6627921275277259, + 1.943570693424962, + 0.4319364830723235, + 1.255200508583198, + -1.3657297839233669, + -0.01336199019878765, + -0.3393069462873377, + 0.10884694751201296, + -0.8166380331391334, + -0.3706884901668285, + 0.573737624799559, + 0.8648002171342688, + -0.03232429463294054, + 0.5307183066940667, + -0.8630078781729218, + 0.6021562424905322, + -0.8453025461956412, + 1.82658398151856, + 0.5628498724685331, + 0.2439306013273806, + -0.07366398394503038, + 1.2700613621652759, + -0.6369663434120979, + 0.7182610420412423, + -1.2498111948276467, + 0.11839550585559121, + 0.6732055094305546, + 0.8361026952993864, + 0.060824890674981126, + -0.6875845134943815, + -0.7631834108102469, + 0.9760045290068342, + -0.7764955885467189, + 1.7629330045268279, + -0.35504472682736055, + 0.21253553542145115, + -1.0526008881405187, + -0.16656267846511705, + 0.4292512537198382, + 1.0647348591723098, + -0.7031255947553559, + 0.0727618965092522, + 0.30766412122713804, + 1.3224517746204512, + -1.3962031323934576, + 0.3726355110942625, + -0.3339235817057772, + -0.16595633888865696, + 0.040995004058702995, + 1.898240920028487, + 1.0099610840499202, + -1.4778201749113464, + 1.3871987975207034, + -1.02136875967816, + 0.6860076952082405, + -0.43995605826138, + 0.114212679611265, + -0.07694607628767526, + 0.6033915854223885, + 0.2905546252040811, + -0.3512251163848229, + 0.423545494803197, + 0.34984439544936924, + 0.6192799614319217, + -0.12000594508590802, + 0.37638481745100344, + -0.07361656093975473, + 0.1594170731485163, + -1.1853656577148597, + 0.425584449110405, + -0.3040611152369916, + -0.3007913065612787, + -0.5506951605940389, + 0.36012535729977696, + 0.5901188186270819, + -0.11362249148148017, + 0.06574730424045437, + -0.03934096367717334, + -0.20957253183767113, + 0.5284990952855384, + 0.3279675860082166, + 0.6123515557513174, + 0.5328126070541334, + 0.32746026978870224, + -0.7679342108202849, + -0.032789070637326424, + 0.2027203335513228, + -0.4361103109207129, + -1.611606821198701, + 0.22342666473445216, + 0.7715215838425302, + -0.14594054046534827, + -1.241469779704863, + -0.09136738204622909, + -1.2382383210774293, + 1.5411714827503882, + -0.6890339762040898, + 2.1760693406110523, + 0.6673206230235658, + -0.1835656432982903, + 1.0249648846349062, + -0.21393771509994333, + 0.3847415726849628, + ] +).reshape((2, 6, 16)) + +_GOLDEN_CALL_DENSE_B = np.array( + [ + -0.048464349831002264, + 0.12578814100862684, + -0.5023540889336157, + 0.10237461495825154, + 0.9474244082687194, + 0.6389273077690177, + -0.3750687372458193, + 0.09538392495089612, + 0.2057911146693594, + -0.32761325570956845, + -0.8137945977073056, + 0.058104558306177854, + 0.6878261146120073, + 0.12799479249653034, + -0.13265034854586666, + -0.24167290953944234, + 0.011129002879949498, + -1.2673340880838346, + 0.21422018979845414, + 0.2376226530727919, + 1.0463862207314913, + 0.19242216959898456, + -0.49840404243215153, + -0.6985139237420458, + 0.6156564351614421, + -0.24904728503654996, + 1.7735818867149378, + 0.9364818091277682, + -0.10997363505719217, + 0.5498063302746796, + 0.24344482865921666, + -0.18352683867020564, + -0.5207070541287808, + -0.0005716822472628322, + 0.43636630351918926, + 0.5333429921068772, + -0.263417441450637, + -0.026220016900765974, + 0.3660861465815261, + 0.0543480783146649, + 0.9332986431668936, + -0.3892111895434158, + 0.9892755313166647, + -0.3914065926656158, + 0.628957621557961, + -1.2955038838182076, + 0.8427745190841428, + -0.11360380972926276, + -0.2606945801433464, + -0.2262084619051, + 0.18969605423739705, + 0.2615142011677701, + 0.7136884709209146, + 0.5763061230504399, + -0.7738108644417824, + 0.49581954056878286, + 0.21465715021208298, + 0.2416302457608278, + -0.6015687398551549, + 0.32567380127953915, + 0.16315710937654773, + -0.02032387835522999, + -0.6476449833730196, + -0.1611926657032835, + ] +).reshape((1, 4, 16)) + + +def test_call_dense_golden() -> None: + """Pin the dense ``call`` outputs across the Task 7 graph-adapter reroute. + + Two fixtures exercise the two dense entry contracts on an edge-sensitive + (jittered zero-init) descriptor so the pin is not vacuous: + + * Fixture A -- ``mapping=None`` gas-phase local indices (``nall == nloc``). + * Fixture B -- a real periodic ~4-atom box with ghosts folded through an + explicit extended->local ``mapping`` (``nall > nloc``), so the + ghost-source scatter is actually exercised. + + The reroute (dense ``call`` -> ``graph_from_dense_quartet`` -> + ``_run_graph``) must preserve these values bit-for-bit within fp64 + scatter-reassociation tolerance. + """ + # Fixture A: mapping=None local indices + ddA = make_message_sensitive_descriptor() + coord, atype, nlist = make_inputs() + nf, nloc = atype.shape + outA = np.asarray(ddA.call(coord.reshape(nf, -1), atype, nlist)[0]) + np.testing.assert_allclose(outA, _GOLDEN_CALL_DENSE_A, rtol=1e-10, atol=1e-12) + + # Fixture B: real periodic ghosts + explicit mapping + ddB = make_message_sensitive_descriptor() + box = np.eye(3, dtype=np.float64)[None] * 6.0 + rng = np.random.default_rng(3) + coord_b = rng.uniform(0.0, 6.0, size=(1, 4, 3)) + atype_b = np.array([[0, 1, 2, 0]], dtype=np.int64) + ext_coord, ext_atype, mapping, nlist_b = extend_input_and_build_neighbor_list( + coord_b, + atype_b, + ddB.get_rcut(), + ddB.get_sel(), + mixed_types=ddB.mixed_types(), + box=box, + ) + assert ext_atype.shape[1] > coord_b.shape[1] # ghosts present + outB = np.asarray(ddB.call(ext_coord, ext_atype, nlist_b, mapping=mapping)[0]) + np.testing.assert_allclose(outB, _GOLDEN_CALL_DENSE_B, rtol=1e-10, atol=1e-12) + + +def test_dense_call_comm_dict_raises() -> None: + # The dense lower has no comm implementation; the dense adapter is the + # one owner of that rejection. + dd = make_descriptor() + coord, atype, nlist = make_inputs() + nf, nloc = atype.shape + with pytest.raises(NotImplementedError, match="dense"): + dd.call(coord.reshape(nf, -1), atype, nlist, comm_dict={"dummy": None}) + + +def test_call_graph_comm_dict_reaches_leaf_stub() -> None: + # The graph trunk now threads comm_dict; in the pure-dpmodel backend the + # per-block exchange leaf is the guard (mirrors dpa2's + # _exchange_ghosts_graph base). ``make_descriptor()`` leaves + # ``use_env_seed`` at its class default (True), so ``_block_comm`` + # forwards comm_dict starting at block 0 already (block 0 is only + # skipped when ``use_env_seed=False``) -- the leaf raise fires on the + # first block regardless of ``n_blocks``. ``n_blocks=2`` is not load + # bearing for reaching the leaf here; it just matches the shared + # ``make_descriptor()`` fixture used across this file. + dd = make_descriptor() + coord, atype, nlist = make_inputs() + graph = make_graph_from_nlist(coord, nlist) + fake_comm = dict.fromkeys( + ( + "send_list", + "send_proc", + "recv_proc", + "send_num", + "recv_num", + "communicator", + "nlocal", + "nghost", + ) + ) + with pytest.raises(NotImplementedError, match="dpmodel backend"): + dd.call_graph(graph, atype.reshape(-1), comm_dict=fake_comm) diff --git a/source/tests/common/dpmodel/test_dpa4_native_spin_model.py b/source/tests/common/dpmodel/test_dpa4_native_spin_model.py new file mode 100644 index 0000000000..86ee99309f --- /dev/null +++ b/source/tests/common/dpmodel/test_dpa4_native_spin_model.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""dpmodel tests for :class:`NativeSpinEnergyModel` (model-level native spin).""" + +import copy + +import numpy as np +import pytest + +from deepmd.dpmodel.model.base_model import ( + BaseModel, +) +from deepmd.dpmodel.model.model import ( + get_model, +) +from deepmd.dpmodel.model.native_spin_model import ( + NativeSpinEnergyModel, +) + +from ...dpa4_fixtures import ( + jitter_zero_arrays, +) + +NATIVE_SPIN_CONFIG = { + "type_map": ["Ni", "O"], + "descriptor": { + "type": "dpa4", + "rcut": 4.0, + "sel": 8, + "channels": 16, + "n_radial": 8, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "precision": "float64", + "seed": 7, + "random_gamma": False, + }, + "fitting_net": {"type": "dpa4_ener", "neuron": [8, 8]}, + "spin": {"use_spin": [True, False], "scheme": "native"}, +} + + +def _jittered_model(seed: int): + """Build the native-spin model then jitter its zero-init residual arrays. + + A fresh DPA4 zero-initializes several residual projections, so it is + architecturally edge/spin-independent (see + ``dpa4_fixtures.jitter_zero_arrays`` docstring); serialize -> jitter -> + deserialize makes the sensitivity/round-trip tests non-vacuous. + """ + model = get_model(NATIVE_SPIN_CONFIG) + data = model.serialize() + data = jitter_zero_arrays(data, np.random.default_rng(seed)) + return BaseModel.deserialize(data) + + +class TestDPA4NativeSpinModel: + def setup_method(self): + self.model = _jittered_model(seed=11) + rng = np.random.default_rng(5) + self.nf, self.nloc = 1, 6 + self.coord = rng.uniform(0.5, 5.5, size=(self.nf, self.nloc, 3)) + self.atype = np.array([[0, 0, 1, 0, 1, 1]], dtype=np.int64) + self.spin = rng.normal(size=(self.nf, self.nloc, 3)) + self.box = 8.0 * np.eye(3, dtype=np.float64)[None] + + def test_call_returns_energy_and_mask_mag(self): + out = self.model.call(self.coord, self.atype, self.spin, box=self.box) + assert out["energy"].shape == (self.nf, 1) + assert out["atom_energy"].shape == (self.nf, self.nloc, 1) + np.testing.assert_array_equal( + out["mask_mag"][..., 0], self.atype == 0 + ) # use_spin=[True, False] + + def test_spin_sensitivity(self): + out0 = self.model.call(self.coord, self.atype, self.spin, box=self.box) + out1 = self.model.call(self.coord, self.atype, 2.0 * self.spin, box=self.box) + assert not np.allclose(out0["energy"], out1["energy"]) + + def test_serialize_roundtrip(self): + data = self.model.serialize() + assert data["type"] == "native_spin" + model2 = BaseModel.deserialize(data) + out0 = self.model.call(self.coord, self.atype, self.spin, box=self.box) + out1 = model2.call(self.coord, self.atype, self.spin, box=self.box) + np.testing.assert_allclose(out0["energy"], out1["energy"], rtol=1e-12) + + @pytest.mark.parametrize( + "legacy_type", + [ + "dpa4_native_spin", # descriptor-specific pre-rename wire type + "sezm_native_spin", # pt backend's own wire type (pt payload layout differs) + ], + ) + def test_legacy_wire_types_fail_fast(self, legacy_type): + # NativeSpinEnergyModel is descriptor-agnostic; mapping it to a + # descriptor-specific wire string would be confusing, so the ONLY + # registered type is "native_spin". Legacy strings raise the + # registry's unknown-type error instead of silently dispatching. + data = self.model.serialize() + data["type"] = legacy_type + with pytest.raises(RuntimeError, match=legacy_type): + BaseModel.deserialize(data) + + def test_dense_route_spin_raises(self): + # The model IS the standard model (is-a): call_common is its own. + with pytest.raises(NotImplementedError, match="NeighborGraph"): + self.model.call_common( + self.coord, + self.atype, + self.box, + spin=self.spin, + neighbor_graph_method="legacy", + ) + + def test_deepspin_scheme_with_dpa4_raises(self): + cfg = {**NATIVE_SPIN_CONFIG, "spin": {"use_spin": [True, False]}} + with pytest.raises(NotImplementedError): + get_model(cfg) + + def test_non_native_spin_descriptor_raises(self): + # The gate is the ``supports_native_spin()`` capability, not a + # descriptor-type list. + cfg = copy.deepcopy(NATIVE_SPIN_CONFIG) + cfg["descriptor"] = { + "type": "se_e2_a", + "rcut": 4.0, + "rcut_smth": 3.5, + "sel": [8, 8], + } + with pytest.raises(NotImplementedError, match="native spin"): + get_model(cfg) + + def test_unrelated_construction_error_propagates(self): + # A bogus fitting kwarg must surface as the REAL TypeError, not be + # masked as a native-spin capability failure (review 3644847676). + cfg = copy.deepcopy(NATIVE_SPIN_CONFIG) + cfg["fitting_net"] = {**cfg["fitting_net"], "bogus_option": 1} + with pytest.raises(TypeError, match="bogus_option"): + get_model(cfg) + + def test_add_chg_spin_ebd_combined_builds_and_conditions(self): + # Combined public configuration (review 3638047227): charge-spin + # FiLM together with native spin, as in pt's SeZMNativeSpinModel. + cfg = copy.deepcopy(NATIVE_SPIN_CONFIG) + cfg["descriptor"]["add_chg_spin_ebd"] = True + model = get_model(cfg) + assert model.has_chg_spin_ebd() + assert model.has_spin() + # Jitter: fresh DPA4 zero-init is architecturally input-independent. + data = model.serialize() + data = jitter_zero_arrays(data, np.random.default_rng(13)) + model = BaseModel.deserialize(data) + dim_cs = model.get_dim_chg_spin() + assert dim_cs > 0 + # charge_spin is CATEGORICAL: ChargeSpinEmbedding casts the frame + # (charge, spin) pair to int64 lookup indices, so only integer-valued + # changes condition the model (0.5 would truncate to 0 == baseline). + cs0 = np.zeros((self.nf, dim_cs), dtype=np.float64) + cs1 = np.array([[1.0, 2.0]], dtype=np.float64) + out_base = model.call( + self.coord, self.atype, self.spin, box=self.box, charge_spin=cs0 + ) + out_cs = model.call( + self.coord, self.atype, self.spin, box=self.box, charge_spin=cs1 + ) + # charge_spin conditions the energy... + assert not np.allclose(out_base["energy"], out_cs["energy"]) + # ...and spin still conditions it in the SAME combined model. + out_spin = model.call( + self.coord, self.atype, 2.0 * self.spin, box=self.box, charge_spin=cs0 + ) + assert not np.allclose(out_base["energy"], out_spin["energy"]) + + def test_translated_output_def_has_spin_keys(self): + out_def = self.model.translated_output_def() + assert "mask_mag" in out_def + assert "force_mag" in out_def + assert "energy" in out_def + assert "force" in out_def + + +class TestNativeSpinConfigForms: + """``spin.use_spin`` index/symbol forms and ``allow_missing_label``. + + The public schema accepts a per-type boolean list, a list of magnetic + type indices, or a list of element symbols (expanded against + ``type_map`` by ``normalize_spin_use_spin``); ``allow_missing_label`` + must be forwarded into the constructed :class:`Spin`. + """ + + @pytest.mark.parametrize( + ("use_spin_form", "expected"), + [ + (["Ni"], [True, False]), # element-symbol form + ([0], [True, False]), # type-index form + ([0, 1], [True, True]), # multiple type indices + ([True, False], [True, False]), # canonical boolean passthrough + ], + ) + def test_use_spin_forms(self, use_spin_form, expected): + config = copy.deepcopy(NATIVE_SPIN_CONFIG) + config["spin"] = {"use_spin": use_spin_form, "scheme": "native"} + model = get_model(config) + assert model.spin.use_spin.tolist() == expected + # The descriptor consumes the SAME normalized boolean list. + descriptor = model.atomic_model.descriptor + assert [bool(flag) for flag in descriptor.use_spin] == expected + + def test_use_spin_unknown_symbol_raises(self): + config = copy.deepcopy(NATIVE_SPIN_CONFIG) + config["spin"] = {"use_spin": ["Fe"], "scheme": "native"} + with pytest.raises(ValueError, match="absent from type_map"): + get_model(config) + + def test_allow_missing_label_forwarded(self): + config = copy.deepcopy(NATIVE_SPIN_CONFIG) + config["spin"]["allow_missing_label"] = True + model = get_model(config) + assert model.spin.allow_missing_label is True + + def test_allow_missing_label_default_false(self): + model = get_model(copy.deepcopy(NATIVE_SPIN_CONFIG)) + assert model.spin.allow_missing_label is False + + +class TestNativeSpinModelRegistryDispatch: + """Wire-type + registry contract of ``make_native_spin_model`` classes. + + Review 3638137290 (PR #5884): dispatch goes through each backend's + plugin registry (backend-aware), not a hard-coded branch in + ``BaseBaseModel.deserialize``; the serialized shape is the make_model + FLAT dict + ``spin`` field, not the legacy nested wrapper shape. + """ + + def _inputs(self): + rng = np.random.default_rng(5) + coord = rng.uniform(0.5, 5.5, size=(1, 6, 3)) + atype = np.array([[0, 0, 1, 0, 1, 1]], dtype=np.int64) + spin = rng.normal(size=(1, 6, 3)) + box = 8.0 * np.eye(3, dtype=np.float64)[None] + return coord, atype, spin, box + + def test_serialize_emits_native_spin_flat_shape(self): + model = get_model(copy.deepcopy(NATIVE_SPIN_CONFIG)) + assert type(model) is NativeSpinEnergyModel + data = model.serialize() + assert data["type"] == "native_spin" + assert "spin" in data + assert "backbone_model" not in data # make_model flat shape, not nested + + def test_basemodel_deserialize_dispatches_via_registry(self): + model = get_model(copy.deepcopy(NATIVE_SPIN_CONFIG)) + m2 = BaseModel.deserialize(model.serialize()) + assert type(m2) is NativeSpinEnergyModel + coord, atype, spin, box = self._inputs() + e1 = model.call(coord, atype, spin, box=box)["energy"] + e2 = m2.call(coord, atype, spin, box=box)["energy"] + np.testing.assert_allclose(e1, e2, rtol=1e-12) diff --git a/source/tests/common/dpmodel/test_dpa4_sparse_edges.py b/source/tests/common/dpmodel/test_dpa4_sparse_edges.py deleted file mode 100644 index edf6107b54..0000000000 --- a/source/tests/common/dpmodel/test_dpa4_sparse_edges.py +++ /dev/null @@ -1,167 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -"""Layout-agnostic edge-aggregation regression test for the dpmodel DPA4. - -The standard ``call`` path consumes a padded edge cache (``E = n_nodes * nnei`` -with ``dst == repeat(arange(n_nodes), nnei)``), while ``call_with_edges`` -consumes an arbitrary sparse edge list. The destination-wise aggregations -(geometric initial embedding, environment initial embedding, and the attention -softmax) are scatter reductions over ``dst``, so both layouts must yield the -same descriptor for the same physical edges. These tests feed the sparse path -the padded path's valid edges -- once in row-major order and once permuted into -an arbitrary order with a non-uniform per-node degree -- and assert the two -descriptors agree. They are the regression guard for the scatter-by-``dst`` -aggregation. -""" - -import numpy as np - -from deepmd.dpmodel.descriptor.dpa4 import ( - DescrptDPA4, -) - - -def build_neighbor_list_np(coord, rcut, nnei): - """Build a padded, distance-sorted gas-phase neighbor list (no PBC). - - Parameters - ---------- - coord - Coordinates with shape (nf, nloc, 3). - rcut - Cutoff radius. - nnei - Number of neighbor slots; pads with -1. - - Returns - ------- - np.ndarray - Neighbor list with shape (nf, nloc, nnei) holding local indices. - """ - nf, nloc, _ = coord.shape - nlist = -np.ones((nf, nloc, nnei), dtype=np.int64) - for f in range(nf): - dist = np.linalg.norm(coord[f][:, None, :] - coord[f][None, :, :], axis=-1) - for i in range(nloc): - neighbors = [ - (dist[i, j], j) for j in range(nloc) if j != i and dist[i, j] < rcut - ] - neighbors.sort() - for slot, (_, j) in enumerate(neighbors[:nnei]): - nlist[f, i, slot] = j - return nlist - - -def build_sparse_edges_from_nlist(coord, nlist): - """Extract the valid physical edges of a padded neighbor list. - - The padded layout keeps one slot per neighbor (``-1`` marks padding). The - sparse contract for :meth:`DescrptDPA4.call_with_edges` is one explicit edge - per kept slot, indexing the flattened frame-major node axis - (``node = f * nloc + i``). The edge vector points from the center toward the - neighbor, matching the padded path's ``r_j - r_i``. - - Parameters - ---------- - coord - Coordinates with shape (nf, nloc, 3). - nlist - Neighbor list with shape (nf, nloc, nnei); -1 marks padding. - - Returns - ------- - tuple[np.ndarray, np.ndarray] - ``edge_index`` with shape (2, E) (rows are src, dst) and ``edge_vec`` - with shape (E, 3), aligned on the same edge axis in row-major - ``(frame, center, slot)`` order. - """ - nf, nloc, nnei = nlist.shape - src, dst, vec = [], [], [] - for f in range(nf): - for i in range(nloc): - for s in range(nnei): - j = int(nlist[f, i, s]) - if j < 0: - continue - src.append(f * nloc + j) - dst.append(f * nloc + i) - vec.append(coord[f, j] - coord[f, i]) - edge_index = np.asarray([src, dst], dtype=np.int64) # (2, E) - edge_vec = np.asarray(vec, dtype=np.float64) # (E, 3) - return edge_index, edge_vec - - -def make_descriptor() -> DescrptDPA4: - return DescrptDPA4( - ntypes=3, - sel=8, - rcut=4.0, - channels=16, - n_radial=8, - lmax=2, - mmax=1, - n_blocks=2, - precision="float64", - seed=7, - random_gamma=False, - ) - - -def make_inputs(seed=7, nf=2, nloc=6, rcut=4.0, nnei=8, ntypes=3): - rng = np.random.default_rng(seed) - coord = rng.uniform(0.0, 3.5, size=(nf, nloc, 3)) - atype = rng.integers(0, ntypes, size=(nf, nloc)) - nlist = build_neighbor_list_np(coord, rcut, nnei) - return coord, atype, nlist - - -def _run_sparse(dd, coord, atype, edge_index, edge_vec): - nf = atype.shape[0] - edge_mask = np.ones(edge_index.shape[1], dtype=bool) - return np.asarray( - dd.call_with_edges( - coord_ext=coord, - atype_ext=atype, - edge_index=edge_index, - edge_vec=edge_vec, - edge_mask=edge_mask, - )[0] - ) - - -def test_sparse_edges_match_padded_rowmajor() -> None: - # Row-major sparse edges reproduce the padded scatter order exactly: the - # masked padding slots of the padded path contribute zero, so dropping them - # leaves the destination accumulation order unchanged. - dd = make_descriptor() - coord, atype, nlist = make_inputs() - nf, nloc = atype.shape - - out_pad = np.asarray(dd.call(coord.reshape(nf, -1), atype, nlist)[0]) - edge_index, edge_vec = build_sparse_edges_from_nlist(coord, nlist) - out_sparse = _run_sparse(dd, coord, atype, edge_index, edge_vec) - - assert out_sparse.shape == out_pad.shape == (nf, nloc, dd.get_dim_out()) - assert np.isfinite(out_sparse).all() - np.testing.assert_allclose(out_sparse, out_pad, rtol=1e-10, atol=1e-12) - - -def test_sparse_edges_match_padded_permuted() -> None: - # Permuted sparse edges exercise an arbitrary (non-row-major) ``dst`` order - # and a non-uniform per-node degree. The destination scatter reductions are - # order-agnostic, so the descriptor must still match the padded path within - # float64 reassociation tolerance. - dd = make_descriptor() - coord, atype, nlist = make_inputs() - nf, nloc = atype.shape - - out_pad = np.asarray(dd.call(coord.reshape(nf, -1), atype, nlist)[0]) - - edge_index, edge_vec = build_sparse_edges_from_nlist(coord, nlist) - perm = np.random.default_rng(31).permutation(edge_index.shape[1]) - edge_index = edge_index[:, perm] - edge_vec = edge_vec[perm] - out_sparse = _run_sparse(dd, coord, atype, edge_index, edge_vec) - - assert out_sparse.shape == out_pad.shape == (nf, nloc, dd.get_dim_out()) - assert np.isfinite(out_sparse).all() - np.testing.assert_allclose(out_sparse, out_pad, rtol=1e-10, atol=1e-12) diff --git a/source/tests/common/dpmodel/test_inter_potential.py b/source/tests/common/dpmodel/test_inter_potential.py new file mode 100644 index 0000000000..20075900d6 --- /dev/null +++ b/source/tests/common/dpmodel/test_inter_potential.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""dpmodel ``InterPotential`` (analytical ZBL bridging term) unit tests. + +Ports the pt reference values from +``source/tests/pt/model/test_sezm_model.py::TestInterPotential``: the exact +universal-ZBL formula is reproduced in-test and the half-split per-edge +scatter must sum back to the full analytic pair energy. +""" + +import math + +import numpy as np +import pytest + +from deepmd.dpmodel.atomic_model.inter_potential import ( + InterPotential, +) + +_A_BOHR = 0.5291772109 +_KE = 14.3996 +_A_COEFF = (0.18175, 0.50986, 0.28022, 0.028171) +_B_COEFF = (3.1998, 0.94229, 0.4029, 0.20162) + + +def _analytic_zbl(r: float, zi: float, zj: float) -> float: + a = 0.88534 * _A_BOHR / (zi**0.23 + zj**0.23) + x = r / a + phi = sum( + a_k * math.exp(-b_k * x) for a_k, b_k in zip(_A_COEFF, _B_COEFF, strict=True) + ) + return _KE * zi * zj / r * phi + + +def _two_atom_inputs(r: float): + """Two atoms at distance r with BOTH directed edges (symmetric list).""" + edge_vec = np.array([[r, 0.0, 0.0], [-r, 0.0, 0.0]], dtype=np.float64) + edge_index = np.array([[0, 1], [1, 0]], dtype=np.int64) # src, dst + edge_mask = np.array([True, True]) + return edge_vec, edge_index, edge_mask + + +@pytest.mark.parametrize( + ("type_map", "atypes", "zi", "zj"), + [ + (["O"], [0, 0], 8.0, 8.0), # O-O pair + (["O", "H"], [0, 1], 8.0, 1.0), # O-H pair + ], +) +def test_zbl_known_value(type_map, atypes, zi, zj): + r = 0.8 + pot = InterPotential(type_map=type_map) + edge_vec, edge_index, edge_mask = _two_atom_inputs(r) + out = pot.call( + edge_vec, + edge_index, + np.asarray(atypes, dtype=np.int64), + edge_mask, + n_node=2, + ) + assert out.shape == (1, 2, 1) + # Half per directed edge -> the total is the full analytic pair energy. + np.testing.assert_allclose( + float(np.sum(out)), _analytic_zbl(r, zi, zj), rtol=0, atol=1e-5 + ) + + +def test_virtual_types_masked(): + # real_type_count=1: type 1 is a virtual/placeholder type; its edges + # contribute zero, and only real-real edges survive. + pot = InterPotential(type_map=["O"]) + r = 0.9 + edge_vec = np.array( + [[r, 0, 0], [-r, 0, 0], [0, r, 0], [0, -r, 0]], dtype=np.float64 + ) + edge_index = np.array([[0, 1, 0, 2], [1, 0, 2, 0]], dtype=np.int64) + edge_mask = np.ones(4, dtype=bool) + atypes = np.array([0, 0, 1], dtype=np.int64) # atom 2 is virtual + out = pot.call(edge_vec, edge_index, atypes, edge_mask, 3, real_type_count=1) + np.testing.assert_allclose( + float(np.sum(out)), _analytic_zbl(r, 8.0, 8.0), atol=1e-5 + ) + assert float(out[0, 2, 0]) == 0.0 + + +def test_edge_mask_zeroes_edges(): + pot = InterPotential(type_map=["O"]) + edge_vec, edge_index, _ = _two_atom_inputs(0.8) + out = pot.call( + edge_vec, + edge_index, + np.zeros(2, dtype=np.int64), + np.array([False, False]), + n_node=2, + ) + np.testing.assert_array_equal(np.asarray(out), 0.0) + + +def test_unknown_element_raises(): + with pytest.raises(ValueError, match="Unknown element symbol"): + InterPotential(type_map=["O", "Xx"]) + + +def test_unknown_mode_raises(): + with pytest.raises(ValueError, match="Unknown InterPotential mode"): + InterPotential(type_map=["O"], mode="lj") + + +def test_torch_namespace_smoke_and_gradient(): + """Torch inputs match numpy at 1e-12 and edge_vec gradients exist.""" + import torch + + pot = InterPotential(type_map=["O", "H"]) + edge_vec_np, edge_index, edge_mask = _two_atom_inputs(0.8) + atypes = np.array([0, 1], dtype=np.int64) + ref = np.asarray(pot.call(edge_vec_np, edge_index, atypes, edge_mask, 2)) + + ev = torch.tensor(edge_vec_np, dtype=torch.float64, requires_grad=True) + out = pot.call( + ev, + torch.tensor(edge_index), + torch.tensor(atypes), + torch.tensor(edge_mask), + 2, + ) + np.testing.assert_allclose(out.detach().numpy(), ref, rtol=1e-12) + grad = torch.autograd.grad(out.sum(), ev)[0] + assert torch.isfinite(grad).all() + assert grad.abs().max().item() > 0.0 diff --git a/source/tests/common/dpmodel/test_linear_atomic_model.py b/source/tests/common/dpmodel/test_linear_atomic_model.py index 79e1bdab75..5630fd5c61 100644 --- a/source/tests/common/dpmodel/test_linear_atomic_model.py +++ b/source/tests/common/dpmodel/test_linear_atomic_model.py @@ -37,6 +37,16 @@ def mixed_types(self) -> bool: def get_type_map(self) -> list[str]: return self.type_map + def get_intensive(self) -> bool: + # part of the atomic-model interface the composition validates + return False + + def get_dim_fparam(self) -> int: + return 0 + + def get_dim_aparam(self) -> int: + return 0 + def change_type_map( self, type_map: list[str], diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index 1109c21131..242ef0c188 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -2238,3 +2238,94 @@ def make_padded(): make_B, make_padded, ) + + +# --------------------------------------------------------------------------- +# Task 11 review fix: EnergySpinLoss flat-label reshape (data-loader shape) +# --------------------------------------------------------------------------- +# The data loader delivers per-atom vector labels FLAT: (nframes, natoms * 3) +# (see ``deepmd/utils/data.py:892``), not the canonical (nframes, natoms, 3) +# shape that ``masked_atom_mean`` requires. ``EnergySpinLoss.call`` bridges +# this via ``xp.reshape(model_dict["force"], (-1, natoms, 3))`` (and the same +# for ``force_mag``). Every test above constructs already-3D labels, so that +# reshape is a no-op there. This test feeds a genuinely FLAT label and pins +# the reshape numerically. + + +class TestDPModelEnerSpinLossFlatLabelReshape: + """Pin the flat-(nf, natoms*3) -> (nf, natoms, 3) label reshape. + + Asserts the loss computed from a flat ``(nf, natoms * 3)`` force / + force_mag label equals the loss computed from the same data pre-reshaped + to ``(nf, natoms, 3)``, with a non-zero force_mag term. This fails (shape + error or wrong value) if the reshape in ``EnergySpinLoss.call`` is + removed. + """ + + def _make_loss(self): + return EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=1.0, + limit_pref_fr=1.0, + start_pref_fm=1.0, + limit_pref_fm=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + ) + + def test_flat_label_matches_3d_label(self): + nf = 2 + force_pred = _rnd(nf, NP, 3) + force_label_3d = _rnd(nf, NP, 3) + force_mag_pred = _rnd(nf, NP, 3) + force_mag_label_3d = _rnd(nf, NP, 3) + mask_mag = _MASK_MAG_PAD_SPIN # [2, NP, 1] + + loss_obj = self._make_loss() + + model_dict = { + "energy": np.zeros((nf, 1), dtype=np.float64), + "force": force_pred, + "force_mag": force_mag_pred, + "mask_mag": mask_mag, + "virial": np.zeros((nf, 9), dtype=np.float64), + } + label_3d = { + "energy": np.zeros((nf, 1), dtype=np.float64), + "force": force_label_3d, + "force_mag": force_mag_label_3d, + "virial": np.zeros((nf, 9), dtype=np.float64), + "find_energy": 0.0, + "find_force": 1.0, + "find_force_mag": 1.0, + "find_virial": 0.0, + } + # The real, motivating shape: the data loader delivers atomic vector + # labels flat -- (nf, natoms * 3) -- not (nf, natoms, 3). + label_flat = dict(label_3d) + label_flat["force"] = force_label_3d.reshape(nf, NP * 3) + label_flat["force_mag"] = force_mag_label_3d.reshape(nf, NP * 3) + + loss_3d, more_loss_3d = loss_obj.call(1.0, NP, model_dict, label_3d) + loss_flat, more_loss_flat = loss_obj.call(1.0, NP, model_dict, label_flat) + + assert np.isclose(float(loss_3d), float(loss_flat), rtol=1e-12, atol=1e-12), ( + f"flat-label loss must equal 3D-label loss: {loss_flat} vs {loss_3d}" + ) + assert float(more_loss_flat["rmse_fm"]) != 0.0, ( + "force_mag loss term must be non-zero for this test to have teeth" + ) + assert np.isclose( + float(more_loss_flat["rmse_fm"]), + float(more_loss_3d["rmse_fm"]), + rtol=1e-12, + atol=1e-12, + ) + assert np.isclose( + float(more_loss_flat["rmse_fr"]), + float(more_loss_3d["rmse_fr"]), + rtol=1e-12, + atol=1e-12, + ) diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py new file mode 100644 index 0000000000..1bca3a2a9e --- /dev/null +++ b/source/tests/common/dpmodel/test_zbl_bridging.py @@ -0,0 +1,631 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""dpmodel ZBL bridging as COMPOSITION (review 3638077323, redesigned). + +``bridging_method: ZBL`` builds a +``LinearEnergyModel(LinearEnergyAtomicModel([dp, InterPotentialAtomicModel], +weights="sum"))`` -- the analytical term is its own atomic model summed +with the learned one, not a flag on it. +""" + +import copy + +import numpy as np +import pytest + +from deepmd.dpmodel.atomic_model.inter_potential import ( + InterPotentialAtomicModel, +) +from deepmd.dpmodel.atomic_model.linear_atomic_model import ( + LinearEnergyAtomicModel, +) +from deepmd.dpmodel.model.base_model import ( + BaseModel, +) +from deepmd.dpmodel.model.dp_linear_model import ( + LinearEnergyModel, +) +from deepmd.dpmodel.model.model import ( + get_model, +) + +ZBL_CONFIG = { + "type_map": ["Ni", "O"], + "descriptor": { + "type": "dpa4", + "rcut": 4.0, + "sel": 8, + "channels": 16, + "n_radial": 8, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "precision": "float64", + "seed": 7, + "random_gamma": False, + }, + "fitting_net": {"type": "dpa4_ener", "neuron": [8, 8]}, + "bridging_method": "ZBL", + "bridging_r_inner": 0.8, + "bridging_r_outer": 1.2, +} + + +def _close_pair_inputs(): + rng = np.random.default_rng(5) + coord = rng.uniform(1.5, 5.5, size=(1, 6, 3)) + coord[0, 1] = coord[0, 0] + np.array([0.9, 0.0, 0.0]) # close Ni-Ni pair + atype = np.array([[0, 0, 1, 0, 1, 1]], dtype=np.int64) + box = 8.0 * np.eye(3, dtype=np.float64)[None] + return coord, atype, box + + +def test_builder_composes_linear_model(): + model = get_model(copy.deepcopy(ZBL_CONFIG)) + assert type(model) is LinearEnergyModel + am = model.atomic_model + assert isinstance(am, LinearEnergyAtomicModel) + assert am.weights == "sum" + kinds = [type(c).__name__ for c in am.models] + assert ( + kinds == ["EnergyAtomicModel", "InterPotentialAtomicModel"] + or kinds[1] == "InterPotentialAtomicModel" + ) + # radii wired to the LEARNED child's descriptor InnerClamp + dp_child = am.models[0] + assert dp_child.descriptor.inner_clamp is not None + assert float(dp_child.descriptor.inner_clamp.r_inner) == 0.8 + + +def test_zbl_child_equals_composition_minus_learned(): + """Composition energy == learned child + analytical child (exact sum).""" + model = get_model(copy.deepcopy(ZBL_CONFIG)) + coord, atype, box = _close_pair_inputs() + e_sum = model.call_common(coord, atype, box=box, neighbor_graph_method="dense")[ + "energy_redu" + ] + dp_child, zbl_child = model.atomic_model.models + # learned child alone through its OWN model wrapper + from deepmd.dpmodel.model.ener_model import ( + EnergyModel, + ) + + m_dp = EnergyModel(atomic_model_=dp_child) + e_dp = m_dp.call_common(coord, atype, box=box, neighbor_graph_method="dense")[ + "energy_redu" + ] + diff = float(np.sum(e_sum - e_dp)) + assert diff > 1e-3, f"ZBL contribution missing or non-positive: {diff:.3e}" + # EXACT analytical check (gas phase, no box: a direct double loop over + # pairs within rcut is the complete reference). + import math + + e_gas_sum = model.call_common(coord, atype, neighbor_graph_method="dense")[ + "energy_redu" + ] + e_gas_dp = m_dp.call_common(coord, atype, neighbor_graph_method="dense")[ + "energy_redu" + ] + z_of = {0: 28.0, 1: 8.0} # Ni, O + total = 0.0 + for i in range(6): + for j in range(i + 1, 6): + r = float(np.linalg.norm(coord[0, i] - coord[0, j])) + if r >= 4.0: + continue + zi, zj = z_of[int(atype[0, i])], z_of[int(atype[0, j])] + a = 0.88534 * 0.5291772109 / (zi**0.23 + zj**0.23) + phi = sum( + ak * math.exp(-bk * (r / a)) + for ak, bk in zip( + (0.18175, 0.50986, 0.28022, 0.028171), + (3.1998, 0.94229, 0.4029, 0.20162), + strict=True, + ) + ) + total += 14.3996 * zi * zj / r * phi + np.testing.assert_allclose( + float(np.sum(e_gas_sum - e_gas_dp)), total, rtol=1e-10, atol=1e-10 + ) + + +def test_zbl_serialize_roundtrip_energy_identical(): + model = get_model(copy.deepcopy(ZBL_CONFIG)) + coord, atype, box = _close_pair_inputs() + data = model.serialize() + # the flat wire type is "linear" -- the SAME string pt/tf write, so a + # composition round-trips across backends + assert data["type"] == "linear" + m2 = BaseModel.deserialize(data) + assert type(m2) is LinearEnergyModel + e1 = model.call_common(coord, atype, box=box, neighbor_graph_method="dense")[ + "energy_redu" + ] + e2 = m2.call_common(coord, atype, box=box, neighbor_graph_method="dense")[ + "energy_redu" + ] + np.testing.assert_allclose(e1, e2, rtol=1e-12) + + +def test_zbl_atomic_dense_route_raises(): + zbl = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + with pytest.raises(NotImplementedError, match="NeighborGraph route only"): + zbl.forward_atomic(None, None, None) + + +def test_inter_potential_supports_graph_lower(): + """The analytical ZBL term is graph-capable (rides the NeighborGraph).""" + zbl = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + assert zbl.uses_graph_lower() is True + + +def test_linear_graph_lower_requires_all_children(): + """A composition is graph-capable iff EVERY child supports the graph lower. + + Regression: a dense-only child (a bare-minimum stand-in for + ``PairTabAtomicModel``, standard DP+ZBL) forces the whole linear model + onto the dense route, even alongside a graph-capable child -- otherwise + the graph route would call ``forward_atomic_graph`` on the dense-only + child, which does not implement it. + """ + + class _GraphChild(InterPotentialAtomicModel): + pass # inherits uses_graph_lower() -> True + + class _DenseOnlyChild(InterPotentialAtomicModel): + def uses_graph_lower(self) -> bool: + return False + + graph_child = _GraphChild(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + dense_child = _DenseOnlyChild(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + + all_graph = LinearEnergyAtomicModel( + [graph_child, _GraphChild(type_map=["Ni", "O"], rcut=4.0, sel=[8])], + type_map=["Ni", "O"], + weights="sum", + ) + assert all_graph.uses_graph_lower() is True + + mixed = LinearEnergyAtomicModel( + [graph_child, dense_child], + type_map=["Ni", "O"], + weights="sum", + ) + assert mixed.uses_graph_lower() is False + + +def test_zbl_atomic_graph_values(): + """Atomic-model wrapper reproduces the kernel's known values.""" + import math + + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + ) + + r = 0.8 + zbl = InterPotentialAtomicModel(type_map=["O"], rcut=4.0, sel=[8]) + graph = NeighborGraph( + n_node=np.array([2], dtype=np.int64), + edge_index=np.array([[0, 1], [1, 0]], dtype=np.int64), + edge_vec=np.array([[r, 0.0, 0.0], [-r, 0.0, 0.0]], dtype=np.float64), + edge_mask=np.ones(2, dtype=bool), + ) + out = zbl.forward_common_atomic_graph(graph, np.zeros(2, dtype=np.int64)) + a = 0.88534 * 0.5291772109 / (8.0**0.23 + 8.0**0.23) + phi = sum( + ak * math.exp(-bk * (r / a)) + for ak, bk in zip( + (0.18175, 0.50986, 0.28022, 0.028171), + (3.1998, 0.94229, 0.4029, 0.20162), + strict=True, + ) + ) + ref = 14.3996 * 64.0 / r * phi + np.testing.assert_allclose(float(np.sum(out["energy"])), ref, atol=1e-5) + + +def _pair_energy(model, natoms=2, r=1.0): + """Total ZBL energy of one pair at distance ``r``, all atoms of type 0.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + ) + + graph = NeighborGraph( + n_node=np.array([natoms], dtype=np.int64), + edge_index=np.array([[0, 1], [1, 0]], dtype=np.int64), + edge_vec=np.array([[r, 0.0, 0.0], [-r, 0.0, 0.0]], dtype=np.float64), + edge_mask=np.ones(2, dtype=bool), + ) + out = model.forward_common_atomic_graph(graph, np.zeros(natoms, dtype=np.int64)) + return float(np.sum(out["energy"])) + + +class TestInterPotentialChangeTypeMap: + """``change_type_map`` must rebuild the ZBL element lookup. + + The generic ``BaseAtomicModel.change_type_map`` only rewrites the public + map and the stat/exclusion state; the nuclear-charge table belongs to + ``InterPotential`` and is rebuilt there (review 3649295675). Without it + the lookup keeps the ORIGINAL elements while ``atype`` values already mean + the new ones -- silently wrong energies, or ``IndexError`` for a longer + map. + """ + + def test_reorder_matches_a_freshly_built_model(self) -> None: + model = InterPotentialAtomicModel(type_map=["H", "O"], rcut=4.0, sel=[8]) + e_hh = _pair_energy(model) + model.change_type_map(["O", "H"]) + fresh = InterPotentialAtomicModel(type_map=["O", "H"], rcut=4.0, sel=[8]) + e_fresh = _pair_energy(fresh) + # anti-vacuity: the two element pairs must be far apart, or a stale + # lookup would be indistinguishable from a rebuilt one + assert abs(e_fresh - e_hh) > 1.0 + np.testing.assert_allclose(_pair_energy(model), e_fresh, rtol=1e-12) + assert list(model.potential.atomic_numbers) == [8.0, 1.0] + assert model.potential.type_map == ["O", "H"] + + def test_added_element_extends_the_lookup(self) -> None: + model = InterPotentialAtomicModel(type_map=["H", "O"], rcut=4.0, sel=[8]) + model.change_type_map(["H", "O", "Ni"]) + assert model.potential.ntypes_real == 3 + assert list(model.potential.atomic_numbers) == [1.0, 8.0, 28.0] + # the new type is now addressable -- a stale (length-2) table raises + # IndexError here + graph_atype = np.full(2, 2, dtype=np.int64) + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + ) + + graph = NeighborGraph( + n_node=np.array([2], dtype=np.int64), + edge_index=np.array([[0, 1], [1, 0]], dtype=np.int64), + edge_vec=np.array([[1.0, 0.0, 0.0], [-1.0, 0.0, 0.0]], dtype=np.float64), + edge_mask=np.ones(2, dtype=bool), + ) + e_nini = float( + np.sum(model.forward_common_atomic_graph(graph, graph_atype)["energy"]) + ) + fresh = InterPotentialAtomicModel(type_map=["Ni"], rcut=4.0, sel=[8]) + np.testing.assert_allclose(e_nini, _pair_energy(fresh), rtol=1e-12) + + def test_dropped_element_shrinks_the_lookup(self) -> None: + model = InterPotentialAtomicModel(type_map=["H", "O", "Ni"], rcut=4.0, sel=[8]) + model.change_type_map(["Ni"]) + assert model.potential.ntypes_real == 1 + assert list(model.potential.atomic_numbers) == [28.0] + + def test_serialize_roundtrip_after_change_type_map(self) -> None: + """Checkpoint continuity: the restored model must predict the same. + + Serialization records the NEW public map, so a stale in-memory lookup + and its deserialized twin disagree -- the restart-time symptom of the + same bug. + """ + from deepmd.dpmodel.atomic_model.base_atomic_model import ( + BaseAtomicModel, + ) + + model = InterPotentialAtomicModel(type_map=["H", "O"], rcut=4.0, sel=[8]) + model.change_type_map(["O", "H"]) + data = model.serialize() + restored = BaseAtomicModel.get_class_by_type(data["type"]).deserialize(data) + assert restored.get_type_map() == ["O", "H"] + np.testing.assert_allclose( + _pair_energy(restored), _pair_energy(model), rtol=1e-12 + ) + + +class TestNativeSpinCapabilityOnAtomicModel: + """``supports_native_spin`` is answered by the ATOMIC MODEL. + + The model layer must not reach into an atomic model for a descriptor to + decide spin eligibility: an analytical term has no descriptor at all, and + a composition has several children. Each atomic model answers from its + own structure, exactly like ``uses_graph_lower``. + """ + + def test_analytical_term_is_not_spin_capable(self) -> None: + zbl = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + # inherits the concrete base default -- no descriptor, no spin input + assert zbl.supports_native_spin() is False + + def test_composition_is_capable_when_any_child_is(self) -> None: + """ANY, not ALL: analytical children accept and ignore ``spin``.""" + learned = get_model( + { + **copy.deepcopy(ZBL_CONFIG), + "spin": {"use_spin": [True, False], "scheme": "native"}, + } + ).atomic_model + assert learned.supports_native_spin() is True + kinds = [type(c).__name__ for c in learned.models] + assert kinds[1] == "InterPotentialAtomicModel", kinds + # ... and the spin-free analytical child alone is not capable + assert learned.models[1].supports_native_spin() is False + + def test_composition_without_a_spin_consumer_is_not_capable(self) -> None: + """No consumer => the magnetic force would be identically zero.""" + zbl_a = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + zbl_b = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + composed = LinearEnergyAtomicModel( + [zbl_a, zbl_b], type_map=["Ni", "O"], weights="sum" + ) + assert composed.supports_native_spin() is False + + +class TestCompositionForwardsConditioningCapabilities: + """A composition must FORWARD every capability its children own. + + ``get_dim_fparam``/``get_dim_aparam`` were forwarded, but the + charge/spin FiLM and default-fparam accessors fell through to + ``BaseAtomicModel``'s ``False``/``0``. That is silently wrong rather + than loudly broken: the eager forward still conditions on + ``charge_spin`` (the learned child consumes it), while the FREEZE reads + these accessors -- so a 0 dropped the charge_spin slot from the exported + ABI and from the metadata the C++ feeder reads, and the artifact + disagreed with its own eager model. + """ + + @staticmethod + def _model(bridging: bool, chg_spin: bool = True): + config = copy.deepcopy(ZBL_CONFIG) + config["descriptor"]["add_chg_spin_ebd"] = chg_spin + if not bridging: + for key in ("bridging_method", "bridging_r_inner", "bridging_r_outer"): + config.pop(key, None) + return get_model(config) + + def test_charge_spin_survives_bridging(self) -> None: + plain = self._model(bridging=False) + bridged = self._model(bridging=True) + # anti-vacuity: the unbridged model must actually declare the input + assert plain.get_dim_chg_spin() > 0 + assert bridged.get_dim_chg_spin() == plain.get_dim_chg_spin() + assert bridged.has_chg_spin_ebd() is True + # ... and the composition really is a composition + assert [type(c).__name__ for c in bridged.atomic_model.models][1] == ( + "InterPotentialAtomicModel" + ) + + def test_no_charge_spin_stays_zero(self) -> None: + """The other branch: nothing is invented when no child declares it.""" + bridged = self._model(bridging=True, chg_spin=False) + assert bridged.get_dim_chg_spin() == 0 + assert bridged.has_chg_spin_ebd() is False + + def test_default_conditioning_accessors_are_forwarded(self) -> None: + """``has_default_*`` must not fall through to the base either.""" + bridged = self._model(bridging=True) + plain = self._model(bridging=False) + assert bridged.has_default_chg_spin() == plain.has_default_chg_spin() + assert bridged.has_default_fparam() == plain.has_default_fparam() + assert bridged.get_default_fparam() == plain.get_default_fparam() + + +class TestCompositionCarriesPairExclusion: + """Model-level ``pair_exclude_types`` must survive the ZBL composition. + + It is a BUILD-time transform: the atomic model only carries the config, + and whoever builds the neighbor graph (the C++ feeder, or the Python + builder) applies it from the exported metadata. ``_collect_metadata`` + reads it off ``model.atomic_model`` -- which for a bridged model is the + ``LinearEnergyAtomicModel``, not the learned child. Building that + composition without forwarding the exclusion therefore dropped it + silently: the eager model still behaved, while the frozen artifact told + its feeder there was nothing to exclude. + """ + + @staticmethod + def _model(bridging: bool, spin: bool = False): + config = copy.deepcopy(ZBL_CONFIG) + config["pair_exclude_types"] = [[0, 1]] + if spin: + config["spin"] = {"use_spin": [True, False], "scheme": "native"} + if not bridging: + for key in ("bridging_method", "bridging_r_inner", "bridging_r_outer"): + config.pop(key, None) + return get_model(config) + + def test_pair_exclusion_survives_bridging(self) -> None: + plain = self._model(bridging=False) + bridged = self._model(bridging=True) + # anti-vacuity: the unbridged model must actually carry it + assert plain.atomic_model.pair_exclude_types == [[0, 1]] + assert bridged.atomic_model.pair_exclude_types == [[0, 1]] + # ... and the bridged one really is the composition, so the value is + # read off the wrapper rather than accidentally off a lone child. + assert [type(c).__name__ for c in bridged.atomic_model.models][1] == ( + "InterPotentialAtomicModel" + ) + + def test_pair_exclusion_survives_native_spin_plus_bridging(self) -> None: + """The three-way stack must not drop it either.""" + assert self._model( + bridging=True, spin=True + ).atomic_model.pair_exclude_types == [[0, 1]] + + def test_no_exclusion_stays_empty(self) -> None: + """The other branch: nothing is invented when none is configured.""" + config = copy.deepcopy(ZBL_CONFIG) + config.pop("pair_exclude_types", None) + assert get_model(config).atomic_model.pair_exclude_types == [] + + +class TestCompositionCarriesAtomExclusion: + """``atom_exclude_types`` must reach the ZBL term too. + + The twin of :class:`TestCompositionCarriesPairExclusion`. Unlike pair + exclusion this one IS applied at runtime, which is why it was initially + (and wrongly) left unforwarded on the theory that forwarding would + double-apply. It does not: the composition's own ``atom_excl`` was + ``None``, the learned child masked only itself, and the analytical child + never heard about the exclusion -- so an excluded atom still collected + its full share of the ZBL energy. Masking to zero is idempotent, so the + child keeping its own copy is harmless. + """ + + @staticmethod + def _model(bridging: bool): + config = copy.deepcopy(ZBL_CONFIG) + config["atom_exclude_types"] = [1] + if not bridging: + for key in ("bridging_method", "bridging_r_inner", "bridging_r_outer"): + config.pop(key, None) + return get_model(config) + + def test_atom_exclusion_reaches_the_composition(self) -> None: + bridged = self._model(bridging=True) + plain = self._model(bridging=False) + # anti-vacuity: the unbridged model must actually carry it + assert plain.atomic_model.atom_exclude_types == [1] + assert bridged.atomic_model.atom_exclude_types == [1] + # the mask is what actually zeroes the analytical child's output + assert bridged.atomic_model.atom_excl is not None + + def test_no_exclusion_stays_empty(self) -> None: + """Nothing is invented when none is configured.""" + config = copy.deepcopy(ZBL_CONFIG) + config.pop("atom_exclude_types", None) + model = get_model(config) + assert model.atomic_model.atom_exclude_types == [] + assert model.atomic_model.atom_excl is None + + +class TestCompositionDefaultsRequireAgreement: + """A parent default is only valid when every ACTIVE child shares it. + + The composition exposes ONE external tensor to all children, so + advertising the first child's default made ``get_additional_data_ + requirement`` mark the input optional and inject that value into every + child -- silently overriding the others' own defaults (reported as an + 8.13e-4 energy change for two learned children defaulting to [0.0] and + [1.0]). Dimension-zero children (an analytical bridging term) are not + consumers and must be ignored, so learned+ZBL still inherits the + learned default. + """ + + class _Fake: + """Learned-child stand-in with its own fparam default.""" + + def __init__(self, dim: int, default) -> None: + self._dim, self._default = dim, default + + def mixed_types(self) -> bool: + return True + + def get_type_map(self) -> list: + return ["Ni", "O"] + + def get_intensive(self) -> bool: + return False + + def get_dim_fparam(self) -> int: + return self._dim + + def get_dim_aparam(self) -> int: + return 0 + + def get_dim_chg_spin(self) -> int: + return 0 + + def has_default_fparam(self) -> bool: + return self._default is not None + + def get_default_fparam(self): + return self._default + + def has_default_chg_spin(self) -> bool: + return False + + def get_default_chg_spin(self): + return None + + def _compose(self, children): + return LinearEnergyAtomicModel(children, type_map=["Ni", "O"], weights="sum") + + def test_consumers_must_agree_on_dimension(self) -> None: + """One shared tensor means one dimension among ACTIVE consumers. + + Rejected at construction, like the intensive/extensive mixture: the + composition feeds every child the same fparam/aparam tensor, so + consumers wanting different widths cannot both be satisfied. + """ + with pytest.raises(ValueError, match="fparam dimension"): + self._compose([self._Fake(3, None), self._Fake(2, None)]) + + def test_dimension_and_default_align_with_the_learned_child(self) -> None: + """Learned + ZBL inherits BOTH the dimension and the default. + + The analytical child consumes neither fparam nor aparam, so it is + not a consumer and must not constrain either. + """ + zbl = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + # anti-vacuity: the analytical child really is a non-consumer + assert zbl.get_dim_fparam() == 0 + assert zbl.get_dim_aparam() == 0 + m = self._compose([self._Fake(3, [0.5, 0.5, 0.5]), zbl]) + assert m.get_dim_fparam() == 3 + assert m.get_default_fparam() == [0.5, 0.5, 0.5] + + def test_differing_defaults_expose_none(self) -> None: + m = self._compose([self._Fake(1, [0.0]), self._Fake(1, [1.0])]) + assert m.has_default_fparam() is False + assert m.get_default_fparam() is None + + def test_matching_defaults_are_exposed(self) -> None: + m = self._compose([self._Fake(1, [1.0]), self._Fake(1, [1.0])]) + assert m.has_default_fparam() is True + assert m.get_default_fparam() == [1.0] + + def test_dimension_zero_child_is_ignored(self) -> None: + """Learned + ZBL must still inherit the learned default.""" + zbl = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + assert zbl.get_dim_fparam() == 0 # anti-vacuity: really a non-consumer + m = self._compose([self._Fake(1, [0.5]), zbl]) + assert m.has_default_fparam() is True + assert m.get_default_fparam() == [0.5] + + def test_active_child_without_default_exposes_none(self) -> None: + m = self._compose([self._Fake(1, [1.0]), self._Fake(1, None)]) + assert m.has_default_fparam() is False + + +class TestCompositionForwardsStatCapabilities: + """``get_intensive`` / ``get_compute_stats_distinguish_types`` aggregate. + + Both silently fell through to ``BaseAtomicModel``'s defaults, so a + bridged model would fit its out-stat bias with the wrong extensivity and + the wrong type-distinguishing rule -- the same class of gap as the + charge-spin and pair-exclusion accessors. + """ + + def test_intensive_mixture_is_rejected_at_construction(self) -> None: + """An intensive/extensive mixture must not be CONSTRUCTIBLE. + + Rejected in ``__init__`` rather than at query time: such a + composition is not physically meaningful, so it should never exist + rather than exist and answer a plausible-looking default. + """ + zbl_a = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + zbl_b = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + # anti-vacuity: matching children compose fine and report their value + assert zbl_a.get_intensive() is False + assert ( + LinearEnergyAtomicModel( + [zbl_a, zbl_b], type_map=["Ni", "O"], weights="sum" + ).get_intensive() + is False + ) + + zbl_b.get_intensive = lambda: True # type: ignore[method-assign] + with pytest.raises(ValueError, match="intensive and extensive"): + LinearEnergyAtomicModel([zbl_a, zbl_b], type_map=["Ni", "O"], weights="sum") + + def test_forwarded_from_children(self) -> None: + bridged = get_model(copy.deepcopy(ZBL_CONFIG)) + children = bridged.atomic_model.models + assert bridged.atomic_model.get_intensive() == all( + c.get_intensive() for c in children + ) + assert bridged.atomic_model.get_compute_stats_distinguish_types() == any( + c.get_compute_stats_distinguish_types() for c in children + ) diff --git a/source/tests/common/test_spin.py b/source/tests/common/test_spin.py index 5249c4f3d1..fa8c91a530 100644 --- a/source/tests/common/test_spin.py +++ b/source/tests/common/test_spin.py @@ -6,11 +6,52 @@ from deepmd.utils.spin import ( Spin, + normalize_spin_use_spin, ) CUR_DIR = os.path.dirname(__file__) +class NormalizeUseSpinTest(unittest.TestCase): + """Unit tests for ``normalize_spin_use_spin`` (pure; all three forms).""" + + def setUp(self) -> None: + self.type_map = ["Ni", "O", "H"] + + def test_boolean_passthrough(self) -> None: + self.assertEqual( + normalize_spin_use_spin([True, False, True], self.type_map), + [True, False, True], + ) + + def test_index_form(self) -> None: + self.assertEqual( + normalize_spin_use_spin([0, 2], self.type_map), + [True, False, True], + ) + + def test_symbol_form(self) -> None: + self.assertEqual( + normalize_spin_use_spin(["Ni", "H"], self.type_map), + [True, False, True], + ) + + def test_empty_list_all_false(self) -> None: + self.assertEqual( + normalize_spin_use_spin([], self.type_map), + [False, False, False], + ) + + def test_unknown_symbol_raises(self) -> None: + with self.assertRaisesRegex(ValueError, "absent from type_map"): + normalize_spin_use_spin(["Fe"], self.type_map) + + def test_pure_no_input_mutation(self) -> None: + use_spin = ["Ni"] + normalize_spin_use_spin(use_spin, self.type_map) + self.assertEqual(use_spin, ["Ni"]) + + class SpinTest(unittest.TestCase): def setUp(self) -> None: type_map_1 = ["H", "O"] diff --git a/source/tests/conftest.py b/source/tests/conftest.py index 9e42da4517..2db5d625c1 100644 --- a/source/tests/conftest.py +++ b/source/tests/conftest.py @@ -36,6 +36,7 @@ "pt_expt/model/test_export_with_comm.py", "pt_expt/model/test_dpa1_graph_lower.py", "pt_expt/model/test_graph_export.py", + "pt_expt/model/test_zbl_bridging.py", "pt_expt/model/test_graph_export_with_comm.py", "pt_expt/utils/test_graph_pt2_metadata.py", "pt_expt/infer/test_deep_eval_metadata_only.py", diff --git a/source/tests/consistent/test_array_api.py b/source/tests/consistent/test_array_api.py index ad40d321d5..05d88f3cc6 100644 --- a/source/tests/consistent/test_array_api.py +++ b/source/tests/consistent/test_array_api.py @@ -10,6 +10,7 @@ xp_scatter_sum, xp_setitem_at, xp_sigmoid, + xp_uniform, ) from deepmd.dpmodel.common import ( to_numpy_array, @@ -370,3 +371,45 @@ def test_tf2_full_rank_mask_consistent_with_ref(self) -> None: tnp.asarray(values_np), ) np.testing.assert_allclose(ref, to_numpy_array(result), atol=1e-10) + + +class TestXpUniform(unittest.TestCase): + """Each backend draws with its own generator, on the reference device.""" + + @unittest.skipUnless(INSTALLED_PT, "PyTorch is not installed") + def test_pt_replays_under_torch_seed(self) -> None: + like = torch.zeros(7, dtype=torch.float64, device=DEVICE) + torch.manual_seed(4321) + a = xp_uniform(like, 512, 0.0, 2.0 * np.pi) + torch.manual_seed(4321) + b = xp_uniform(like, 512, 0.0, 2.0 * np.pi) + torch.manual_seed(1234) + c = xp_uniform(like, 512, 0.0, 2.0 * np.pi) + + assert a.shape == (512,) + assert a.dtype == like.dtype + assert a.device.type == like.device.type + torch.testing.assert_close(a, b, rtol=0.0, atol=0.0) + # anti-vacuity: a constant would satisfy the replay check alone + assert not torch.allclose(a, c) + assert float(a.min()) >= 0.0 + assert float(a.max()) < 2.0 * np.pi + + def test_numpy_fallback_replays_under_the_project_seed(self) -> None: + """The fallback uses deepmd's seeded generator, not a fresh one.""" + from deepmd.utils import random as dp_random + + like = np.zeros(3, dtype=np.float64) + dp_random.seed(777) + a = xp_uniform(like, 64, -1.0, 1.0) + dp_random.seed(777) + b = xp_uniform(like, 64, -1.0, 1.0) + dp_random.seed(778) + c = xp_uniform(like, 64, -1.0, 1.0) + + assert a.shape == (64,) + assert a.dtype == like.dtype + np.testing.assert_array_equal(a, b) + assert not np.allclose(a, c) + assert a.min() >= -1.0 + assert a.max() < 1.0 diff --git a/source/tests/dpa4_fixtures.py b/source/tests/dpa4_fixtures.py new file mode 100644 index 0000000000..a99caf2589 --- /dev/null +++ b/source/tests/dpa4_fixtures.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Shared DPA4 test fixtures used across test roots. + +``jitter_zero_arrays`` is imported by ``source/tests/common/dpmodel``, +``source/tests/pt_expt`` test modules that need a message-sensitive DPA4 +fixture. ``source/tests/infer/gen_dpa4.py`` keeps its own inline copy +(mirror of this module) because it is a standalone script that runs +outside pytest's package machinery. +""" + +import numpy as np + + +def jitter_zero_arrays(node, rng: np.random.Generator): + """Return a copy of a serialized tree with every zero float array jittered. + + DPA4 deliberately zero-initializes several residual output projections + (``SO2Convolution.post_focus_mix``, ``EquivariantFFN.so3_linear_2`` -- + both per-block and the top-level ``output_ffn`` -- see the "Zero- + initialized so residual path starts near-identity" comments in + ``dpa4_nn/so2.py``/``dpa4_nn/ffn.py``) so a freshly constructed, + untrained descriptor is architecturally edge/message independent: its + scalar read-out is exactly the type embedding regardless of geometry, + neighbors, or ``exclude_types``. That makes a bare ``make_descriptor()`` + vacuous for an exclusion anti-vacuity check -- excluding pairs cannot + change an output that never depended on edges. + + This function replaces *every* float array in the serialized weight tree + that is exactly all-zero, wherever it occurs (not only the named residual + projections above) -- it has no notion of which key it is perturbing, it + only inspects array values. Non-zero arrays (e.g. the learned type + embedding) are carried through untouched, and traversal order is a fixed + depth-first walk, so two calls seeded identically produce bit-identical + trees. + + This is a PURE rebuild -- it does not mutate ``node`` (nor anything it + references), so callers must use the return value: + ``data = jitter_zero_arrays(data, rng)``. (Mutating ``node`` in place + tripped CodeQL's ``py/modification-of-default-value`` dataflow.) + + Parameters + ---------- + node : dict, list, np.ndarray, or other + Root (or sub-tree) of a serialized parameter tree. Not mutated. + rng : np.random.Generator + Seeded RNG used to draw the replacement noise. + + Returns + ------- + dict, list, np.ndarray, or other + A new tree of the same shape with zero float arrays replaced; leaves + that are not zero float arrays are returned as-is. + """ + if isinstance(node, dict): + return {key: jitter_zero_arrays(value, rng) for key, value in node.items()} + if isinstance(node, list): + return [jitter_zero_arrays(value, rng) for value in node] + if ( + isinstance(node, np.ndarray) + and node.dtype.kind == "f" + and node.size > 0 + and np.all(node == 0.0) + ): + return rng.normal(0.0, 0.05, size=node.shape).astype(node.dtype) + return node diff --git a/source/tests/infer/gen_dpa4.py b/source/tests/infer/gen_dpa4.py index d09f372e65..3515c19de2 100644 --- a/source/tests/infer/gen_dpa4.py +++ b/source/tests/infer/gen_dpa4.py @@ -120,7 +120,13 @@ def main(): pt2_path = os.path.join(base_dir, "deeppot_dpa4.pt2") print(f"Exporting to {pt2_path} ...") # noqa: T201 - pt_expt_deserialize_to_file(pt2_path, copy.deepcopy(data), do_atomic_virial=True) + # Pinned explicitly: DPA4 is graph-native end-to-end, so + # ``lower_kind="auto"`` now resolves to "graph" (_resolve_lower_kind). + # This dense fixture must keep testing the dense-nlist ABI regardless of + # that default; the graph ABI is exercised separately by Section B below. + pt_expt_deserialize_to_file( + pt2_path, copy.deepcopy(data), do_atomic_virial=True, lower_kind="nlist" + ) pth_path = os.path.join(base_dir, "deeppot_dpa4.pth") print(f"Exporting to {pth_path} ...") # noqa: T201 @@ -234,6 +240,235 @@ def main(): else: print("\n// Skipping .pth verification (file not generated).") # noqa: T201 + # ============================================================ + # Section B: graph .pt2 export (jittered weights, non-vacuous) + # ============================================================ + # DPA4 is graph-native end-to-end: there is no dense-only sub-block to + # toggle off for graph-eligibility (unlike DPA2's use_three_body), so + # the SAME config as Section A is graph-eligible and, before the pin + # above, ``lower_kind="auto"`` already resolved to "graph" for it (see + # deepmd/pt_expt/utils/serialization.py:_resolve_lower_kind). + # + # Skip the whole graph section under LeakSanitizer. The C++ memleak + # matrix runs these gen scripts with ``LD_PRELOAD=liblsan`` (the + # sanitizer-instrumented deepmd op .so requires the LSAN runtime; see + # source/install/test_cc_local.sh). Evaluating the AOTInductor-compiled + # graph .pt2's BACKWARD (forces) under that runtime INTERMITTENTLY + # segfaults -- an AOTI-compiled-code vs LeakSanitizer allocator + # incompatibility, NOT a graph-code bug (see gen_dpa2.py's identical + # section for the full rationale -- dpa2's repformer trips the same + # issue). The dense DPA4 .pt2 (Section A) is unaffected and still + # generated. The C++ dpa4_graph_pytorch_pt2 row GTEST_SKIPs when this + # artifact is absent (skip_if_artifact_missing). + # + # Detection is via the explicit DP_GEN_UNDER_SANITIZER flag set by + # test_cc_local.sh next to the preload: sniffing LD_PRELOAD here does + # NOT reliably work -- the LSAN runtime removes its own entry from the + # process environment during startup on some platforms. The LD_PRELOAD + # check is kept only as a belt-and-braces fallback for manual + # invocations where the runtime leaves the variable intact. + if ( + os.environ.get("DP_GEN_UNDER_SANITIZER", "") == "lsan" + or "lsan" in os.environ.get("LD_PRELOAD", "").lower() + ): + # remove any graph artifacts left by a previous non-LSAN run of a + # REUSED workspace: skipping regeneration alone leaves them present, + # and the C++ tests' skip_if_artifact_missing would then execute + # them under LSAN and hit the very crash this branch avoids + for name in ( + "deeppot_dpa4_graph_nlist_ref.pt2", + "deeppot_dpa4_graph.pt2", + "deeppot_dpa4_graph.expected", + ): + stale = os.path.join(base_dir, name) + if os.path.exists(stale): + os.remove(stale) + print( # noqa: T201 + "\n// Skipping DPA4 graph section under LeakSanitizer " + "(AOTInductor .pt2 backward is incompatible with the LSAN runtime; " + "covered by the non-memleak C++/LAMMPS matrix)." + ) + print("\nDone!") # noqa: T201 + return + + print("\n---- Building graph DPA4 (jittered weights) ----") # noqa: T201 + + # ---- B.1 Build a fresh DPA4 model; jitter zero-init residuals ---- + # A freshly built DPA4 zero-initializes several residual output + # projections (see step 2b above and the ``jitter_zero_arrays`` + # docstring in source/tests/dpa4_fixtures.py), so its output is + # architecturally edge-independent until those branches are perturbed + # away from exactly zero. Section A already does this via in-place + # torch-parameter replacement (step 2b); this section instead follows + # the dict-level ``jitter_zero_arrays`` pattern used by + # test_dpa4_call_graph.py / test_dpa4_graph_lower.py, for consistency + # with the rest of the DPA4 graph test suite. Inlined here (rather than + # imported from source/tests/dpa4_fixtures.py) because gen_dpa4.py is a + # standalone script run outside pytest's package machinery -- importing + # a source/tests/... module from it would need ad hoc sys.path / + # package surgery for no real benefit. + # Mirror of source/tests/dpa4_fixtures.py:jitter_zero_arrays -- keep in sync. + def _jitter_zero_arrays(node, rng: np.random.Generator): + # Mirror of source/tests/dpa4_fixtures.py:jitter_zero_arrays -- keep in + # sync. PURE rebuild (returns a new tree, does not mutate ``node``) to + # avoid CodeQL's py/modification-of-default-value dataflow; behavior + # (RNG draws, shapes, dtype) is bit-identical. + if isinstance(node, dict): + return {k: _jitter_zero_arrays(v, rng) for k, v in node.items()} + if isinstance(node, list): + return [_jitter_zero_arrays(v, rng) for v in node] + if ( + isinstance(node, np.ndarray) + and node.dtype.kind == "f" + and node.size > 0 + and np.all(node == 0.0) + ): + return rng.normal(0.0, 0.05, size=node.shape).astype(node.dtype) + return node + + model_g = get_model(copy.deepcopy(config)) + model_g.to("cpu") + model_g.eval() + model_dict_g = model_g.serialize() + model_dict_g = _jitter_zero_arrays(model_dict_g, np.random.default_rng(20240615)) + + data_g = { + "model": copy.deepcopy(model_dict_g), + "model_def_script": config, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + + # ---- B.2 Independent cross-check via nlist .pt2 (dense-quartet) ---- + # Like gen_dpa2.py's Section B.2: deeppot_dpa4_graph.expected is NOT + # copied from the nlist artifact (see B.5), since DPA4's graph path + # assigns each edge's force/virial contribution fully to the source atom + # (edge_force_virial full-to-src), a different (equally valid) + # decomposition than the dense per-atom one -- only the SUM agrees. The + # nlist .pt2 is instead used here as the independent gen-time oracle: + # atomic energies, forces and the TOTAL virial of the graph .pt2 must + # match it (checked in B.4) or generation aborts. + # + # The nlist .pt2 is PERSISTED (deeppot_dpa4_graph_nlist_ref.pt2), reused + # directly by the LAMMPS graph-vs-nlist-ref test and by the DeepEval + # graph parity test, both exercised on the SAME jittered weights as the + # graph model, so at non-binding sel the two paths must agree. + nlist_ref_pt2 = os.path.join(base_dir, "deeppot_dpa4_graph_nlist_ref.pt2") + print(f"Exporting reference nlist .pt2 to {nlist_ref_pt2} ...") # noqa: T201 + pt_expt_deserialize_to_file( + nlist_ref_pt2, + copy.deepcopy(data_g), + do_atomic_virial=True, + lower_kind="nlist", # independent: dense nlist, NOT graph + ) + dp_nlist_ref = DeepPot(nlist_ref_pt2) + + # PBC reference from nlist path + e_r1, f_r1, v_r1, ae_r1, av_r1 = dp_nlist_ref.eval(coord, box, atype, atomic=True) + # NoPBC reference from nlist path + e_rnp, f_rnp, v_rnp, ae_rnp, av_rnp = dp_nlist_ref.eval( + coord, None, atype, atomic=True + ) + + print(f"Nlist ref PBC energy: {e_r1[0, 0]:.18e}") # noqa: T201 + print(f"Nlist ref NoPBC energy: {e_rnp[0, 0]:.18e}") # noqa: T201 + max_ref_force_pbc = float(np.max(np.abs(f_r1))) + max_ref_force_nopbc = float(np.max(np.abs(f_rnp))) + print(f"Nlist ref PBC max |force|: {max_ref_force_pbc:.6e}") # noqa: T201 + print(f"Nlist ref NoPBC max |force|: {max_ref_force_nopbc:.6e}") # noqa: T201 + # Anti-vacuity guard: a fresh DPA4 is edge-independent (zero-init + # residual projections), so a broken (or accidentally skipped) jitter + # above would silently produce a degenerate, geometry-insensitive + # fixture. ``not (x >= th)`` (rather than ``x < th``) so NaN forces -- + # e.g. from an inductor SIMD miscompile of the AOTI artifact -- fail the + # check instead of slipping through. + if ( + not (max_ref_force_pbc > 1e-6) + or not (max_ref_force_nopbc > 1e-6) + or not (np.all(np.isfinite(f_r1)) and np.all(np.isfinite(f_rnp))) + ): + raise RuntimeError( + f"BLOCKED: graph DPA4 nlist-ref forces are degenerate or " + f"non-finite (PBC max={max_ref_force_pbc:.2e}, " + f"NoPBC max={max_ref_force_nopbc:.2e}); the zero-init-residual " + f"jitter may have failed to perturb the descriptor, or the AOTI " + f"compile is broken (known inductor CPU-SIMD bug; workaround: " + f"torch._inductor.config.cpp.simdlen = 1)." + ) + + # ---- B.3 Export graph-form .pt2 (SAME jittered weights) ---- + graph_pt2_path = os.path.join(base_dir, "deeppot_dpa4_graph.pt2") + print(f"Exporting to {graph_pt2_path} (lower_kind='graph') ...") # noqa: T201 + # has_message_passing_across_ranks() is True -> the graph export + # auto-embeds model/extra/forward_lower_with_comm.pt2 (multi-rank + # LAMMPS). + pt_expt_deserialize_to_file( + graph_pt2_path, + copy.deepcopy(data_g), + do_atomic_virial=True, + lower_kind="graph", + ) + print("Graph .pt2 export done.") # noqa: T201 + + # ---- B.4 Cross-check: graph .pt2 vs independent nlist reference ---- + # Both use the SAME weights; at non-binding sel the math is equivalent. + # Atomic energies, forces and the TOTAL virial must agree. The per-atom + # virial is deliberately NOT compared: see B.2. + dp_graph = DeepPot(graph_pt2_path) + + e_g1, f_g1, v_g1, ae_g1, av_g1 = dp_graph.eval(coord, box, atype, atomic=True) + e_gnp, f_gnp, v_gnp, ae_gnp, av_gnp = dp_graph.eval(coord, None, atype, atomic=True) + + cross_tol = 1e-8 + for label, (f_g, ae_g, v_g), (f_r, ae_r, v_r) in ( + ("PBC", (f_g1, ae_g1, v_g1), (f_r1, ae_r1, v_r1)), + ("NoPBC", (f_gnp, ae_gnp, v_gnp), (f_rnp, ae_rnp, v_rnp)), + ): + f_diff = float(np.max(np.abs(f_g[0] - f_r[0]))) + ae_diff = float(np.max(np.abs(ae_g[0] - ae_r[0]))) + v_diff = float(np.max(np.abs(v_g[0] - v_r[0]))) + print( # noqa: T201 + f"Graph .pt2 vs nlist ref {label}: ae {ae_diff:.2e}, " + f"f {f_diff:.2e}, total-virial {v_diff:.2e}" + ) + # NaN-safe: NaN fails ``<=`` + if not (f_diff <= cross_tol and ae_diff <= cross_tol and v_diff <= cross_tol): + raise RuntimeError( + f"BLOCKED: graph .pt2 {label} differs from nlist reference " + f"(ae {ae_diff:.2e}, f {f_diff:.2e}, v {v_diff:.2e}; " + f"threshold {cross_tol:.0e})." + ) + + # ---- B.5 Write sidecar reference file from the graph .pt2 eval ---- + # Self-referential like the dense fixtures' .expected (the C++ gtest is a + # regression test of the C++ inference path against the Python eval of + # the same artifact); independence from the graph path is enforced above + # in B.4. Sourcing e/f/v from the nlist artifact instead would break the + # per-atom virial comparison (convention, see B.2) and sit at the + # gtest's 1e-10 double tolerance for energies/forces (cross-path noise + # is only checked to 1e-8 here). + graph_ref_path = os.path.join(base_dir, "deeppot_dpa4_graph.expected") + write_expected_ref( + graph_ref_path, + sections={ + "pbc": { + "expected_e": ae_g1[0, :, 0], + "expected_f": f_g1[0], + "expected_v": av_g1[0], + }, + "nopbc": { + "expected_e": ae_gnp[0, :, 0], + "expected_f": f_gnp[0], + "expected_v": av_gnp[0], + }, + }, + source_script="source/tests/infer/gen_dpa4.py", + ) + print(f"Wrote {graph_ref_path}") # noqa: T201 + + print("\nAll graph sanity checks passed.") # noqa: T201 + print("\nDone!") # noqa: T201 diff --git a/source/tests/infer/gen_dpa4_spin.py b/source/tests/infer/gen_dpa4_spin.py new file mode 100644 index 0000000000..9cbe25c701 --- /dev/null +++ b/source/tests/infer/gen_dpa4_spin.py @@ -0,0 +1,372 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Generate the native-spin DPA4 graph-route .pt2 test models. + +Two archives, IDENTICAL weights, differing only in model-level exclusion: + +- ``deeppot_dpa4_spin_graph.pt2`` -- no exclusion (baseline) +- ``deeppot_dpa4_spin_pairexcl.pt2`` -- ``pair_exclude_types=[[0, 1]]``, with + the descriptor's own ``exclude_types`` left EMPTY. + +The second one is deliberately anti-vacuous for the C++ ingestion seam. Model- +level exclusion is a BUILD-time transform (decision #18/A4) owned by the +neighbor-graph construction, so the .pt2 lower consumes a pre-excluded +``edge_mask`` and never re-applies it; ``DeepSpinPTExpt`` must fold it in at +its build seam (``applyPairExclusion``). Keeping the descriptor exclusion +empty is what makes the fixture load-bearing: the ``type="dpa4"`` model alias +copies ``pair_exclude_types`` into ``descriptor.exclude_types``, which bakes an +equivalent mask INTO the compiled artifact and would mask a dead C++ seam. + +Generation follows ``gen_dpa4.py``'s Section B pattern: the dpmodel is built in-process +from the inline ``NATIVE_SPIN_CONFIG`` below with a fixed weight-init seed, +its zero-initialized residual projections are jittered away from exact zero +with a fixed RNG seed (``jitter_zero_arrays``, imported from +``source/tests/dpa4_fixtures.py``), and the result is frozen directly to the +graph-kind ``.pt2`` -- no intermediate ``.yaml`` is read or written. Both +``get_model``/weight-init and ``np.random.default_rng(seed)`` are +deterministic, so this reproduces byte-identical weights on every machine/CI +run without committing a serialized-weights file to git. + +The native spin scheme has NO dense/nlist lower at all, spin rides the +NeighborGraph lower exclusively (see +``deepmd/pt_expt/model/native_spin_model.py``'s module docstring and +``source/tests/pt_expt/model/test_dpa4_export.py`` Task 6). Without the +jitter, a freshly built DPA4 collapses to a type-embedding-only descriptor +(see ``jitter_zero_arrays``'s docstring): every force AND force_mag would be +identically zero regardless of geometry/spin, making this fixture vacuous +for the C++/LAMMPS consumers (Tasks 9/10). + +Also writes a sidecar ``.expected`` reference file (PBC and NoPbc per-atom +energy/force/force_mag/virial) consumed by the C++ tests, mirroring +``gen_spin.py``'s field convention. +""" + +import copy +import json +import os +import sys +import zipfile + +import numpy as np + +# Ensure the source tree is on the path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..")) +# Ensure source/tests is on the path for dpa4_fixtures (this script runs +# standalone, outside pytest's package machinery, so the usual `from +# ...dpa4_fixtures import ...` relative import used by the test suite does +# not apply here). +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from dpa4_fixtures import ( + jitter_zero_arrays, +) +from gen_common import ( + derive_pair_exclude_pt2, + ensure_inductor_compiler, + load_custom_ops, + write_expected_ref, +) + +# Small fp64 DPA4/SeZM native-spin config. Mirrors ``NATIVE_SPIN_CONFIG`` in +# source/tests/common/dpmodel/test_dpa4_native_spin_model.py and +# source/tests/pt_expt/model/test_dpa4_native_spin.py (the config exercised +# by Tasks 1-7's dpmodel/pt_expt/export tests): channels 16, n_radial 8, +# lmax 2, mmax 1, n_blocks 2 -- large enough to exercise the SO(2)/SO(3) + +# attention + spin-embedding paths, small enough to keep the AOTInductor +# compile bounded. ``use_spin=[True, False]``: type 0 ("Ni") carries a +# magnetic moment, type 1 ("O") does not. ``scheme="native"`` selects the +# NeighborGraph-only spin route (no virtual atoms, unlike the deepspin +# ``spin_ener`` scheme used by ``gen_spin.py``). +NATIVE_SPIN_CONFIG = { + "type_map": ["Ni", "O"], + "descriptor": { + "type": "dpa4", + "rcut": 4.0, + "sel": 8, + "channels": 16, + "n_radial": 8, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "precision": "float64", + "seed": 7, + "random_gamma": False, + }, + "fitting_net": { + "type": "dpa4_ener", + "neuron": [8, 8], + "precision": "float64", + "seed": 7, + }, + "spin": {"use_spin": [True, False], "scheme": "native"}, +} + +# Fixed seed for jittering the zero-initialized residual projections away +# from exact zero (see ``jitter_zero_arrays``'s docstring and module +# docstring above). Kept as the value the fixture was originally generated +# with, for continuity of the fixed 6-atom reference numbers below. +_JITTER_SEED = 20260720 + +# Model-level exclusion for the second archive: drop every Ni-O pair. Applied +# ONLY at the model level -- ``NATIVE_SPIN_CONFIG["descriptor"]`` carries no +# ``exclude_types``, so nothing inside the compiled artifact reproduces it (see +# the module docstring). +_PAIR_EXCLUDE_TYPES = [[0, 1]] + + +def _build_model_dict() -> dict: + """Build the native-spin dpmodel from config+seed and jitter in place.""" + from deepmd.dpmodel.model.model import ( + get_model, + ) + + model = get_model(copy.deepcopy(NATIVE_SPIN_CONFIG)) + model_dict = model.serialize() + model_dict = jitter_zero_arrays(model_dict, np.random.default_rng(_JITTER_SEED)) + return model_dict + + +# Fixed 6-atom system (3 Ni, spin-active; 3 O, non-magnetic). Coordinates +# and spins reused verbatim from +# source/tests/pt_expt/model/test_dpa4_export.py's +# ``_SPIN_EVAL_{COORDS,CELL,SPINS}`` / ``_SPIN_EVAL_ATYPES`` -- a system +# already validated (Task 7) to yield a non-degenerate ``force_mag`` with +# this same architecture (rcut=4.0, sel=8). Spin is deliberately NOT +# pre-masked by type: the model's own descriptor gating must zero the +# non-spin (type 1 / O) rows internally. +_NATOMS = 6 +_ATYPES = np.array([0, 0, 0, 1, 1, 1], dtype=np.int32) # Ni, Ni, Ni, O, O, O +_COORDS = np.array( + [ + [1.0, 1.0, 1.0], + [3.2, 1.4, 1.1], + [1.3, 1.8, 1.0], + [0.4, 1.2, 1.6], + [3.6, 2.0, 1.3], + [3.4, 0.7, 1.7], + ], + dtype=np.float64, +).reshape(1, _NATOMS, 3) +_CELL = (np.eye(3, dtype=np.float64) * 6.0).reshape(1, 9) +_SPINS = np.array( + [ + [0.11, 0.05, -0.02], + [-0.07, 0.09, 0.03], + [0.02, -0.06, 0.08], + [0.01, -0.01, 0.02], + [-0.02, 0.03, -0.01], + [0.015, 0.02, -0.03], + ], + dtype=np.float64, +).reshape(1, _NATOMS, 3) + + +def _check_metadata(pt2_path: str, expected_exclude: list) -> None: + """Assert the frozen archive's metadata, including the exclusion field.""" + with zipfile.ZipFile(pt2_path) as zf: + md = json.loads(zf.read("model/extra/metadata.json").decode("utf-8")) + print("\n// metadata:") # noqa: T201 + print( # noqa: T201 + json.dumps( + { + k: md[k] + for k in ( + "type_map", + "lower_input_kind", + "is_spin", + "has_comm_artifact", + "has_message_passing", + "ntypes_spin", + "use_spin", + "pair_exclude_types", + "output_keys", + ) + if k in md + }, + indent=2, + ) + ) + assert md["type_map"] == NATIVE_SPIN_CONFIG["type_map"] + assert md["lower_input_kind"] == "graph", ( + f"expected native-spin DPA4 to freeze to the graph lower, got " + f"{md.get('lower_input_kind')!r}" + ) + assert md["is_spin"] is True + # Native spin rides the with-comm artifact on the graph lower, so the + # archive carries the nested forward_lower_with_comm.pt2 for the C++ + # multi-rank path (a SECOND inductor compile -- this generator is the + # slowest of the set for that reason). + assert md["has_comm_artifact"] is True + assert md["has_message_passing"] is True + assert md["use_spin"] == [True, False] + # The exclusion travels as METADATA -- work still owed by the feeder, NOT + # compiled into the artifact. + assert md.get("pair_exclude_types", []) == expected_exclude, ( + f"{pt2_path}: metadata pair_exclude_types = " + f"{md.get('pair_exclude_types')!r}, expected {expected_exclude!r}" + ) + for key in ("atom_energy", "energy", "force", "force_mag", "virial"): + assert key in md["output_keys"] + + +def _eval_and_write_ref(pt2_path: str, ref_path: str) -> float: + """Evaluate one archive (PBC + NoPbc) and write its ``.expected`` sidecar. + + Returns the PBC total energy so the caller can assert that the excluded + archive is not numerically identical to the baseline. + """ + from deepmd.infer import ( + DeepPot, + ) + + dp = DeepPot(pt2_path) + assert dp.has_spin + + e1, f1, v1, ae1, av1, fm1, _mm1 = dp.eval( + _COORDS, _CELL, _ATYPES, atomic=True, spin=_SPINS + ) + print(f"\n// {pt2_path} PBC total energy: {e1[0, 0]:.18e}") # noqa: T201 + + e_np, f_np, v_np, ae_np, av_np, fm_np, _mm_np = dp.eval( + _COORDS, None, _ATYPES, atomic=True, spin=_SPINS + ) + print(f"// {pt2_path} NoPbc total energy: {e_np[0, 0]:.18e}") # noqa: T201 + + spin_mask = _ATYPES == 0 # Ni carries spin; O does not + for label, e, f, fm in ( + ("PBC", e1, f1, fm1), + ("NoPbc", e_np, f_np, fm_np), + ): + assert np.all(np.isfinite(e)), f"{label}: non-finite energy" + assert np.all(np.isfinite(f)), f"{label}: non-finite force" + assert np.all(np.isfinite(fm)), f"{label}: non-finite force_mag" + + fm_flat = fm.reshape(_NATOMS, 3) + fm_spin_max = float(np.max(np.abs(fm_flat[spin_mask]))) + fm_nospin_max = float(np.max(np.abs(fm_flat[~spin_mask]))) + print( # noqa: T201 + f"// {label} max |force_mag| on spin atoms: {fm_spin_max:.6e}" + ) + print( # noqa: T201 + f"// {label} max |force_mag| on non-spin atoms: {fm_nospin_max:.6e}" + ) + # Anti-vacuity: a fresh (non-jittered) DPA4 zero-initializes its + # residual projections, making force_mag identically zero on the + # spin-carrying atoms too -- would silently produce a degenerate + # fixture (see the jitter docstring above). + assert fm_spin_max > 1e-6, ( + f"{label}: expected non-trivial force_mag on spin-active (Ni) " + f"atoms; got max |force_mag| = {fm_spin_max:.3e} (jitter not " + f"effective -- this fixture would be vacuous)." + ) + # The non-spin (O) rows must be exactly gated to zero by the + # model's own type mask -- not merely small. + assert fm_nospin_max == 0.0, ( + f"{label}: expected force_mag to be EXACTLY zero on non-spin " + f"(O) atoms; got max |force_mag| = {fm_nospin_max:.3e}." + ) + + write_expected_ref( + ref_path, + sections={ + "pbc": { + "expected_e": ae1[0, :, 0], + "expected_f": f1[0], + "expected_fm": fm1[0], + "expected_tot_v": v1[0], + "expected_atom_v": av1[0], + }, + "nopbc": { + "expected_e": ae_np[0, :, 0], + "expected_f": f_np[0], + "expected_fm": fm_np[0], + "expected_tot_v": v_np[0], + "expected_atom_v": av_np[0], + }, + }, + source_script="source/tests/infer/gen_dpa4_spin.py", + ) + print(f"Wrote {ref_path}") # noqa: T201 + return float(e1[0, 0]) + + +def main(): + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file as pt_expt_deserialize_to_file, + ) + + ensure_inductor_compiler() + load_custom_ops() + + base_dir = os.path.dirname(__file__) + base_pt2 = os.path.join(base_dir, "deeppot_dpa4_spin_graph.pt2") + excl_pt2 = os.path.join(base_dir, "deeppot_dpa4_spin_pairexcl.pt2") + + # ---- 1. Build the jittered dpmodel dict from config+seed ---- + model_dict = _build_model_dict() + # Negative contract: nothing inside the artifact may reproduce a + # model-level exclusion, or a dead external seam would go unnoticed. The + # generic (no ``type: dpa4``) config above is what keeps this empty. + descrpt_excl = model_dict["descriptor"].get("exclude_types") or [] + assert not descrpt_excl, ( + f"descriptor exclude_types must stay EMPTY for the derived " + f"pair-exclusion fixture to be load-bearing; got {descrpt_excl!r}" + ) + data = { + "model": model_dict, + "model_def_script": NATIVE_SPIN_CONFIG, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + + # ---- 2. Freeze directly to graph-kind .pt2 (the ONLY inductor compile) ---- + # Native-spin DPA4 has NO dense/nlist lower at all (spin rides the + # NeighborGraph lower exclusively -- see the module docstring above), so + # ``lower_kind="auto"`` resolves to "graph" for this model + # (``_resolve_lower_kind``); the virtual-atom ``spin_ener`` scheme would + # instead hard-stop at "nlist". Pinned explicitly here for clarity. + print(f"Exporting to {base_pt2} (lower_kind='graph') ...") # noqa: T201 + pt_expt_deserialize_to_file( + base_pt2, data, do_atomic_virial=True, lower_kind="graph" + ) + print("Export done.") # noqa: T201 + _check_metadata(base_pt2, []) + e_base = _eval_and_write_ref( + base_pt2, os.path.join(base_dir, "deeppot_dpa4_spin_graph.expected") + ) + + # ---- 3. Derive the pair-excluded variant -- NO second compile ---- + # Model-level exclusion is not baked into the exported graph, so patching + # the two JSON blobs that carry the list yields an archive whose compiled + # AOTI artifact is byte-identical to the baseline (see + # gen_common.derive_pair_exclude_pt2). That identity is exactly what makes + # the C++ regression sharp: the two archives can only differ through the + # ingestion seam. + print(f"\nDeriving {excl_pt2} from {base_pt2} ...") # noqa: T201 + derive_pair_exclude_pt2(base_pt2, excl_pt2, _PAIR_EXCLUDE_TYPES) + _check_metadata(excl_pt2, _PAIR_EXCLUDE_TYPES) + e_excl = _eval_and_write_ref( + excl_pt2, os.path.join(base_dir, "deeppot_dpa4_spin_pairexcl.expected") + ) + + # ---- 4. Anti-vacuity for the exclusion itself ---- + # Dropping every Ni-O pair must move the energy. Equal energies would mean + # the exclusion never reached the graph build, making the C++ regression + # that consumes these references pass for the wrong reason. + print( # noqa: T201 + f"\n// baseline PBC energy: {e_base:.18e}\n" + f"// excluded PBC energy: {e_excl:.18e}\n" + f"// delta: {abs(e_excl - e_base):.6e}" + ) + assert abs(e_excl - e_base) > 1e-6, ( + f"pair_exclude_types={_PAIR_EXCLUDE_TYPES} left the energy unchanged " + f"({e_base:.18e} vs {e_excl:.18e}); the exclusion is not reaching the " + f"neighbor-graph build, so both fixtures would be vacuous." + ) + + print("\nDone!") # noqa: T201 + + +if __name__ == "__main__": + main() diff --git a/source/tests/infer/gen_dpa4_spin_chgspin.py b/source/tests/infer/gen_dpa4_spin_chgspin.py new file mode 100644 index 0000000000..502767c6ac --- /dev/null +++ b/source/tests/infer/gen_dpa4_spin_chgspin.py @@ -0,0 +1,392 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Generate the COMBINED native-spin + charge-spin FiLM DPA4 .pt2 fixture. + +One archive, ``deeppot_dpa4_spin_chgspin.pt2``: a native-spin (``scheme= +"native"``) DPA4 whose descriptor ALSO carries ``add_chg_spin_ebd=True`` +(``get_dim_chg_spin() == 2``) and a stored ``default_chg_spin``. That +combination is a supported public configuration (see +``source/tests/pt_expt/model/test_dpa4_native_spin.py``'s +``TestCombinedChargeSpin*``); ``charge_spin`` rides the CONDITIONAL slot-13 +tail of the graph-spin ABI, after the native ``spin`` input at slot 10. + +Why this fixture exists +----------------------- +The C++/LAMMPS spin inference path grew a runtime ``charge_spin`` argument +(``DeepSpin::compute(..., charge_spin)``, ``DP_DeepSpinCompute3``, +``pair_deepspin``'s keyword). Before this fixture nothing demonstrated that a +runtime ``charge_spin`` actually reaches the model: every existing spin +fixture has ``dim_chg_spin == 0``, so the whole argument was inert and a dead +seam would have been invisible. The sibling non-spin fixture +(``gen_chg_spin.py`` -> ``chg_spin.pt2``, DPA3) covers ``DeepPot`` only; the +``DeepSpin`` ingestion seam is a SEPARATE code path +(``DeepSpinPTExpt::compute``, two overloads) and needs its own model. + +Generation mirrors ``gen_dpa4_spin.py`` exactly: the dpmodel is built +in-process from ``NATIVE_SPIN_CONFIG`` (imported from that script -- this +fixture is that model plus the charge-spin FiLM, and nothing else) with a +fixed weight-init seed, its zero-initialized residual projections are +jittered away from exact zero with a fixed RNG seed +(``jitter_zero_arrays``), and the result is frozen directly to the +graph-kind ``.pt2``. Both ``get_model``/weight-init and +``np.random.default_rng(seed)`` are deterministic, so this reproduces +byte-identical weights on every machine/CI run without committing a +serialized-weights file to git. + +Without the jitter a freshly built DPA4 collapses to a type-embedding-only +descriptor (see ``jitter_zero_arrays``'s docstring): force and force_mag +would be identically zero regardless of geometry/spin, and -- since the +charge-spin FiLM feeds those same zero-initialized residual projections -- +the ``charge_spin`` response would vanish too, making the fixture vacuous +for exactly the property it exists to pin. + +The sidecar ``deeppot_dpa4_spin_chgspin.expected`` carries FOUR sections, +PBC and NoPbc x default and explicit ``charge_spin``: + +- ``pbc_default`` / ``nopbc_default`` -- eval with NO charge_spin, i.e. the + model's stored ``default_chg_spin = [0.0, 1.0]``. This is what an EMPTY + runtime ``charge_spin`` must reproduce in C++ (backward compatibility). +- ``pbc_explicit`` / ``nopbc_explicit`` -- eval with ``charge_spin = + [1.0, 2.0]``. + +``ChargeSpinEmbedding`` is CATEGORICAL (it casts the frame ``(charge, spin)`` +pair to int64 lookup indices: ``charge + 100`` and ``spin``), so the two +probes are integer-valued and land on distinct rows: ``[0.0, 1.0]`` -> +(100, 1), ``[1.0, 2.0]`` -> (101, 2). + +Consumed by ``source/api_cc/tests/test_deepspin_dpa4_chgspin_ptexpt.cc``. +""" + +import copy +import json +import os +import sys +import zipfile + +import numpy as np + +# Ensure the source tree is on the path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..")) +# Ensure source/tests is on the path for dpa4_fixtures (this script runs +# standalone, outside pytest's package machinery, so the usual `from +# ...dpa4_fixtures import ...` relative import used by the test suite does +# not apply here). +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from dpa4_fixtures import ( + jitter_zero_arrays, +) +from gen_common import ( + ensure_inductor_compiler, + load_custom_ops, + write_expected_ref, +) + +# The base native-spin DPA4 config, shared verbatim with the sibling +# native-spin fixture: this archive is THAT model plus the charge-spin FiLM, +# so importing (rather than re-typing) the config keeps the single difference +# between the two fixtures visible in one place below. +from gen_dpa4_spin import ( + NATIVE_SPIN_CONFIG, +) + +# Stored fallback used whenever the caller passes no charge_spin. Written to +# the .pt2 metadata as ``default_chg_spin`` and read back by +# ``DeepSpinPTExpt::init`` into ``default_chg_spin_``; an EMPTY runtime +# charge_spin must reproduce exactly this. Same value as gen_chg_spin.py's +# non-spin DPA3 fixture, for cross-fixture consistency. +_DEFAULT_CHG_SPIN = [0.0, 1.0] + +# Explicit runtime probe. Distinct in BOTH components from the default: +# charge idx 101 vs 100, spin idx 2 vs 1 (the embedding is categorical, see +# the module docstring), so neither component alone can explain a response. +_EXPLICIT_CHG_SPIN = [1.0, 2.0] + +CHG_SPIN_CONFIG = copy.deepcopy(NATIVE_SPIN_CONFIG) +CHG_SPIN_CONFIG["descriptor"]["add_chg_spin_ebd"] = True +CHG_SPIN_CONFIG["descriptor"]["default_chg_spin"] = _DEFAULT_CHG_SPIN + +# Fixed seed for jittering the zero-initialized residual projections away +# from exact zero (see ``jitter_zero_arrays``'s docstring and the module +# docstring above). Deliberately NOT gen_dpa4_spin.py's seed: this is a +# different model (extra FiLM weights change the traversal), so sharing a +# seed would only suggest a weight relationship that does not exist. +_JITTER_SEED = 20260726 + + +def _build_model_dict() -> dict: + """Build the combined dpmodel from config+seed and jitter in place.""" + from deepmd.dpmodel.model.model import ( + get_model, + ) + + model = get_model(copy.deepcopy(CHG_SPIN_CONFIG)) + assert model.get_dim_chg_spin() == 2, ( + f"expected the combined native-spin DPA4 to expose dim_chg_spin == 2, " + f"got {model.get_dim_chg_spin()}" + ) + assert model.has_default_chg_spin() + model_dict = model.serialize() + model_dict = jitter_zero_arrays(model_dict, np.random.default_rng(_JITTER_SEED)) + return model_dict + + +# Fixed 6-atom system (3 Ni, spin-active; 3 O, non-magnetic) -- coordinates, +# cell and spins verbatim from gen_dpa4_spin.py, a system already validated +# to yield a non-degenerate ``force_mag`` with this architecture (rcut=4.0, +# sel=8). Spin is deliberately NOT pre-masked by type: the model's own +# descriptor gating must zero the non-spin (type 1 / O) rows internally. +_NATOMS = 6 +_ATYPES = np.array([0, 0, 0, 1, 1, 1], dtype=np.int32) # Ni, Ni, Ni, O, O, O +_COORDS = np.array( + [ + [1.0, 1.0, 1.0], + [3.2, 1.4, 1.1], + [1.3, 1.8, 1.0], + [0.4, 1.2, 1.6], + [3.6, 2.0, 1.3], + [3.4, 0.7, 1.7], + ], + dtype=np.float64, +).reshape(1, _NATOMS, 3) +_CELL = (np.eye(3, dtype=np.float64) * 6.0).reshape(1, 9) +_SPINS = np.array( + [ + [0.11, 0.05, -0.02], + [-0.07, 0.09, 0.03], + [0.02, -0.06, 0.08], + [0.01, -0.01, 0.02], + [-0.02, 0.03, -0.01], + [0.015, 0.02, -0.03], + ], + dtype=np.float64, +).reshape(1, _NATOMS, 3) + + +def _check_metadata(pt2_path: str) -> None: + """Assert the frozen archive's metadata, including the charge-spin slot.""" + with zipfile.ZipFile(pt2_path) as zf: + md = json.loads(zf.read("model/extra/metadata.json").decode("utf-8")) + print("\n// metadata:") # noqa: T201 + print( # noqa: T201 + json.dumps( + { + k: md[k] + for k in ( + "type_map", + "lower_input_kind", + "is_spin", + "has_comm_artifact", + "has_message_passing", + "use_spin", + "dim_chg_spin", + "has_default_chg_spin", + "default_chg_spin", + "output_keys", + ) + if k in md + }, + indent=2, + ) + ) + assert md["type_map"] == CHG_SPIN_CONFIG["type_map"] + assert md["lower_input_kind"] == "graph", ( + f"expected native-spin DPA4 to freeze to the graph lower, got " + f"{md.get('lower_input_kind')!r}" + ) + # Without BOTH of these the artifact does not exercise the feature at all: + # is_spin=False would route the C++ side through DeepPot, and + # dim_chg_spin=0 would make every runtime charge_spin argument inert (which + # is precisely the state every pre-existing spin fixture is in). + assert md["is_spin"] is True, ( + f"{pt2_path}: metadata is_spin = {md.get('is_spin')!r}, expected True; " + f"the DeepSpin charge_spin seam would never be reached." + ) + assert md.get("dim_chg_spin") == 2, ( + f"{pt2_path}: metadata dim_chg_spin = {md.get('dim_chg_spin')!r}, " + f"expected 2; a runtime charge_spin would be silently ignored, making " + f"the C++ regression that consumes this archive vacuous." + ) + assert md.get("has_default_chg_spin") is True, ( + f"{pt2_path}: metadata has_default_chg_spin = " + f"{md.get('has_default_chg_spin')!r}, expected True; the empty-" + f"charge_spin (backward-compatibility) C++ case would throw instead." + ) + stored_default = [float(x) for x in md.get("default_chg_spin", [])] + assert stored_default == _DEFAULT_CHG_SPIN, ( + f"{pt2_path}: metadata default_chg_spin = {stored_default!r}, expected " + f"{_DEFAULT_CHG_SPIN!r}" + ) + # Native spin rides the with-comm artifact on the graph lower, so the + # archive carries the nested forward_lower_with_comm.pt2 for the C++ + # multi-rank path (a SECOND inductor compile). + assert md["has_comm_artifact"] is True + assert md["has_message_passing"] is True + assert md["use_spin"] == [True, False] + for key in ("atom_energy", "energy", "force", "force_mag", "virial"): + assert key in md["output_keys"] + + +def _eval_one(dp, cell, charge_spin, label: str) -> tuple[dict, float]: + """Evaluate one (cell, charge_spin) case; return its ref arrays + energy. + + ``charge_spin=None`` means "pass no charge_spin at all", i.e. exercise the + model's stored ``default_chg_spin`` -- the Python twin of an EMPTY + ``std::vector`` on the C++ side. + """ + kwargs = {} + if charge_spin is not None: + kwargs["charge_spin"] = np.array([charge_spin], dtype=np.float64) + e, f, v, ae, av, fm, _mm = dp.eval( + _COORDS, cell, _ATYPES, atomic=True, spin=_SPINS, **kwargs + ) + print(f"// {label} total energy: {e[0, 0]:.18e}") # noqa: T201 + + assert np.all(np.isfinite(e)), f"{label}: non-finite energy" + assert np.all(np.isfinite(f)), f"{label}: non-finite force" + assert np.all(np.isfinite(fm)), f"{label}: non-finite force_mag" + + spin_mask = _ATYPES == 0 # Ni carries spin; O does not + fm_flat = fm.reshape(_NATOMS, 3) + fm_spin_max = float(np.max(np.abs(fm_flat[spin_mask]))) + fm_nospin_max = float(np.max(np.abs(fm_flat[~spin_mask]))) + print( # noqa: T201 + f"// max |force_mag| spin / non-spin atoms: " + f"{fm_spin_max:.6e} / {fm_nospin_max:.6e}" + ) + # Anti-vacuity: a fresh (non-jittered) DPA4 zero-initializes its residual + # projections, making force_mag identically zero on the spin-carrying + # atoms too (see the jitter docstring above). + assert fm_spin_max > 1e-6, ( + f"{label}: expected non-trivial force_mag on spin-active (Ni) atoms; " + f"got max |force_mag| = {fm_spin_max:.3e} (jitter not effective -- " + f"this fixture would be vacuous)." + ) + # The non-spin (O) rows must be exactly gated to zero by the model's own + # type mask -- not merely small. + assert fm_nospin_max == 0.0, ( + f"{label}: expected force_mag to be EXACTLY zero on non-spin (O) " + f"atoms; got max |force_mag| = {fm_nospin_max:.3e}." + ) + return { + "expected_e": ae[0, :, 0], + "expected_f": f[0], + "expected_fm": fm[0], + "expected_tot_v": v[0], + "expected_atom_v": av[0], + }, float(e[0, 0]) + + +def main(): + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file as pt_expt_deserialize_to_file, + ) + + ensure_inductor_compiler() + load_custom_ops() + + base_dir = os.path.dirname(__file__) + pt2_path = os.path.join(base_dir, "deeppot_dpa4_spin_chgspin.pt2") + ref_path = os.path.join(base_dir, "deeppot_dpa4_spin_chgspin.expected") + + # ---- 1. Build the jittered dpmodel dict from config+seed ---- + model_dict = _build_model_dict() + data = { + "model": model_dict, + "model_def_script": CHG_SPIN_CONFIG, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + + # ---- 2. Freeze directly to graph-kind .pt2 ---- + # Native-spin DPA4 has NO dense/nlist lower at all (spin rides the + # NeighborGraph lower exclusively), so ``lower_kind="auto"`` would resolve + # to "graph" anyway; pinned explicitly here for clarity. + print(f"Exporting to {pt2_path} (lower_kind='graph') ...") # noqa: T201 + pt_expt_deserialize_to_file( + pt2_path, data, do_atomic_virial=True, lower_kind="graph" + ) + print("Export done.") # noqa: T201 + _check_metadata(pt2_path) + + # ---- 3. Evaluate the four reference cases ---- + from deepmd.infer import ( + DeepPot, + ) + + dp = DeepPot(pt2_path) + assert dp.has_spin + dim = dp.deep_eval.get_dim_chg_spin() + assert dim == 2, f"expected dim_chg_spin == 2 from DeepEval, got {dim}" + + print("") # noqa: T201 + pbc_default, e_pbc_default = _eval_one(dp, _CELL, None, "PBC default") + pbc_explicit, e_pbc_explicit = _eval_one( + dp, _CELL, _EXPLICIT_CHG_SPIN, f"PBC explicit {_EXPLICIT_CHG_SPIN}" + ) + nopbc_default, e_nopbc_default = _eval_one(dp, None, None, "NoPbc default") + nopbc_explicit, e_nopbc_explicit = _eval_one( + dp, None, _EXPLICIT_CHG_SPIN, f"NoPbc explicit {_EXPLICIT_CHG_SPIN}" + ) + + # ---- 4. Anti-vacuity: charge_spin must MOVE the output ---- + # Equal energies would mean the charge-spin FiLM never reached the graph + # forward, so the C++ regression consuming these references would pass for + # the wrong reason (it asserts two charge_spin values give two energies). + for label, e_def, e_exp in ( + ("PBC", e_pbc_default, e_pbc_explicit), + ("NoPbc", e_nopbc_default, e_nopbc_explicit), + ): + print( # noqa: T201 + f"\n// {label} default energy: {e_def:.18e}\n" + f"// {label} explicit energy: {e_exp:.18e}\n" + f"// {label} delta: {abs(e_exp - e_def):.6e}" + ) + assert abs(e_exp - e_def) > 1e-6, ( + f"{label}: charge_spin={_EXPLICIT_CHG_SPIN} left the energy " + f"unchanged vs the stored default {_DEFAULT_CHG_SPIN} " + f"({e_def:.18e} vs {e_exp:.18e}); the FiLM conditioning is not " + f"reaching the forward, so this fixture would be vacuous." + ) + + # ---- 5. Pin the "no charge_spin == stored default" contract ---- + # The C++ empty-charge_spin case is checked against the ``*_default`` + # sections, so those sections must really BE the stored default and not + # some third behaviour. Passing the default value explicitly has to + # reproduce them bit-for-bit (same tensor, same code path). + explicit_default, e_explicit_default = _eval_one( + dp, _CELL, _DEFAULT_CHG_SPIN, f"PBC explicit {_DEFAULT_CHG_SPIN} (== default)" + ) + for key, arr in explicit_default.items(): + np.testing.assert_allclose( + arr, + pbc_default[key], + rtol=1e-14, + atol=1e-14, + err_msg=( + f"passing charge_spin={_DEFAULT_CHG_SPIN} explicitly disagrees " + f"with omitting it ({key}); the stored default_chg_spin is not " + f"what the omitted case uses." + ), + ) + assert abs(e_explicit_default - e_pbc_default) < 1e-12 + + # ---- 6. Write the sidecar reference ---- + write_expected_ref( + ref_path, + sections={ + "pbc_default": pbc_default, + "pbc_explicit": pbc_explicit, + "nopbc_default": nopbc_default, + "nopbc_explicit": nopbc_explicit, + }, + source_script="source/tests/infer/gen_dpa4_spin_chgspin.py", + ) + print(f"\nWrote {ref_path}") # noqa: T201 + + print("\nDone!") # noqa: T201 + + +if __name__ == "__main__": + main() diff --git a/source/tests/infer/gen_dpa4_spin_zbl.py b/source/tests/infer/gen_dpa4_spin_zbl.py new file mode 100644 index 0000000000..09d1d4121d --- /dev/null +++ b/source/tests/infer/gen_dpa4_spin_zbl.py @@ -0,0 +1,531 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Generate the COMBINED native-spin + ZBL-bridging DPA4 .pt2 fixture. + +One archive, ``deeppot_dpa4_spin_zbl_graph.pt2``: a native-spin +(``scheme="native"``) DPA4 that is ALSO bridged (``bridging_method: "ZBL"``), +i.e. a ``LinearEnergyModel`` over ``[learned DPA4, InterPotentialAtomicModel]`` +with ``weights="sum"`` wrapped by the native-spin model class. + +Why this fixture exists +----------------------- +Native spin and analytical bridging each had coverage down to C++/LAMMPS +(``gen_dpa4_spin.py`` / ``gen_dpa4_zbl.py``), but their COMBINATION had none +below the pt_expt Python layer: its only export-seam test +(``source/tests/pt_expt/model/test_zbl_bridging.py:: +test_native_spin_with_bridging_graph_freeze_and_deep_eval``) is +``pytest.skip``ped when ``CI=true``, so CI validated the combination through +eager Python alone. The combination is not the conjunction of the two +covered paths: the spin route feeds ``spin`` into a COMPOSITION (the learned +child consumes it, the analytical child accepts and ignores it), and the +freeze must carry ``is_spin`` metadata for a model whose top-level class is +the linear composition -- neither single-feature fixture exercises that. + +Single-rank only, and this fixture PINS that limitation +------------------------------------------------------- +Bridging enables the descriptor's Source Freeze Propagation Gate, whose +per-node ``eta_j = prod_{e: src_e = j} w_e`` folds a node's FULL outgoing-edge +set. Edges exist only for owned centres, so eta is incomplete on every rank +and no with-comm artifact may be exported. ``_check_metadata`` asserts BOTH +``has_comm_artifact is False`` AND that the nested +``model/extra/forward_lower_with_comm.pt2`` entry is absent (mirroring +``gen_dpa4_zbl.py``) -- an artifact appearing there would silently promise a +multi-rank capability the model cannot honour. Note this is the opposite of +the UNbridged native-spin fixture (``gen_dpa4_spin.py``), which does carry +the with-comm twin: bridging is what removes it. + +Generation mirrors ``gen_dpa4_spin_chgspin.py`` (the closest precedent): the +dpmodel is built in-process from ``NATIVE_SPIN_CONFIG`` imported from +``gen_dpa4_spin.py`` -- this fixture is THAT model plus the three bridging +keys, and nothing else -- with a fixed weight-init seed, its zero-initialized +residual projections are jittered away from exact zero with a fixed RNG seed +(``jitter_zero_arrays``), and the result is frozen directly to the graph-kind +``.pt2``. Both ``get_model``/weight-init and ``np.random.default_rng(seed)`` +are deterministic, so this reproduces byte-identical weights on every +machine/CI run without committing a serialized-weights file to git. + +Without the jitter a freshly built DPA4 collapses to a type-embedding-only +descriptor (see ``jitter_zero_arrays``'s docstring): force and force_mag +would be identically zero regardless of geometry/spin, so the LEARNED half of +the composition would contribute nothing and the fixture would only ever test +the analytical term. + +The sidecar ``deeppot_dpa4_spin_zbl_graph.expected`` carries the usual ``pbc`` +/ ``nopbc`` sections (per-atom energy, force, force_mag, total and atomic +virial). Consumed by +``source/api_cc/tests/test_deepspin_dpa4_zbl_ptexpt.cc``. +""" + +import copy +import json +import math +import os +import sys +import zipfile + +import numpy as np + +# Ensure the source tree is on the path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..")) +# Ensure source/tests is on the path for dpa4_fixtures (this script runs +# standalone, outside pytest's package machinery, so the usual `from +# ...dpa4_fixtures import ...` relative import used by the test suite does +# not apply here). +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from dpa4_fixtures import ( + jitter_zero_arrays, +) +from gen_common import ( + ensure_inductor_compiler, + load_custom_ops, + write_expected_ref, +) + +# The base native-spin DPA4 config, shared verbatim with the sibling +# native-spin fixture: this archive is THAT model plus the bridging keys, so +# importing (rather than re-typing) the config keeps the single difference +# between the two fixtures visible in the three lines below. +from gen_dpa4_spin import ( + NATIVE_SPIN_CONFIG, +) + +# Bridging radii, identical to gen_dpa4_zbl.py's: they feed the descriptor's +# InnerClamp AND BridgingSwitch (built together from the same radii) on the +# LEARNED child, and are what makes the composition single-rank only. +_BRIDGING_R_INNER = 0.8 +_BRIDGING_R_OUTER = 1.2 + +SPIN_ZBL_CONFIG = copy.deepcopy(NATIVE_SPIN_CONFIG) +SPIN_ZBL_CONFIG["bridging_method"] = "ZBL" +SPIN_ZBL_CONFIG["bridging_r_inner"] = _BRIDGING_R_INNER +SPIN_ZBL_CONFIG["bridging_r_outer"] = _BRIDGING_R_OUTER + +# Fixed seed for jittering the zero-initialized residual projections away from +# exact zero (see ``jitter_zero_arrays``'s docstring and the module docstring +# above). Deliberately NOT gen_dpa4_spin.py's nor gen_dpa4_spin_chgspin.py's +# seed: this is a different model (the composition adds a second child, so the +# traversal differs), and sharing a seed would only suggest a weight +# relationship that does not exist. +_JITTER_SEED = 20260727 + +# Fixed 6-atom system (3 Ni, spin-active; 3 O, non-magnetic). Coordinates and +# cell verbatim from gen_dpa4_zbl.py -- atoms 0 and 1 sit 0.9 A apart, inside +# ``bridging_r_outer``, so the analytical ZBL term contributes a large, +# unmistakable repulsion instead of a numerical afterthought (and BOTH members +# of that close pair are spin-carrying Ni, so the ZBL and spin channels act on +# the same atoms). Spins verbatim from gen_dpa4_spin.py. Spin is deliberately +# NOT pre-masked by type: the model's own descriptor gating must zero the +# non-spin (type 1 / O) rows internally. +_NATOMS = 6 +_ATYPES = np.array([0, 0, 0, 1, 1, 1], dtype=np.int32) # Ni, Ni, Ni, O, O, O +_COORDS = np.array( + [ + [1.0, 1.0, 1.0], + [1.9, 1.0, 1.0], + [1.3, 1.8, 1.0], + [0.4, 1.2, 1.6], + [3.6, 2.0, 1.3], + [3.4, 0.7, 1.7], + ], + dtype=np.float64, +).reshape(1, _NATOMS, 3) +_CELL = (np.eye(3, dtype=np.float64) * 6.0).reshape(1, 9) +_SPINS = np.array( + [ + [0.11, 0.05, -0.02], + [-0.07, 0.09, 0.03], + [0.02, -0.06, 0.08], + [0.01, -0.01, 0.02], + [-0.02, 0.03, -0.01], + [0.015, 0.02, -0.03], + ], + dtype=np.float64, +).reshape(1, _NATOMS, 3) + +# ZBL screening parameters, repeated here as an INDEPENDENT reference (they +# mirror deepmd/dpmodel/atomic_model/inter_potential.py, deliberately not +# imported from it: a reference that shares its constants with the code under +# test cannot catch a wrong constant). +_ZBL_A_COEFF = (0.18175, 0.50986, 0.28022, 0.028171) +_ZBL_B_COEFF = (3.1998, 0.94229, 0.4029, 0.20162) +_KE_EV_A = 14.3996 +_A_BOHR = 0.5291772109 +_Z_OF_TYPE = {0: 28.0, 1: 8.0} # type_map ["Ni", "O"] + + +def _analytic_zbl_total(coord: np.ndarray, atype: np.ndarray, rcut: float) -> float: + """Gas-phase ZBL total energy: a direct double loop over pairs < rcut. + + Parameters + ---------- + coord : np.ndarray + (natoms, 3) coordinates in Angstrom. No periodic images are + considered, so this is only a complete reference for a gas-phase + (no-box) evaluation. + atype : np.ndarray + (natoms,) atom types, indexing ``_Z_OF_TYPE``. + rcut : float + Cutoff radius; pairs at or beyond it contribute nothing. + + Returns + ------- + float + The summed ZBL pair energy in eV. + """ + zs = [_Z_OF_TYPE[int(t)] for t in atype] + total = 0.0 + natoms = len(zs) + for ii in range(natoms): + for jj in range(ii + 1, natoms): + r = float(np.linalg.norm(coord[ii] - coord[jj])) + if r >= rcut: + continue + a_screen = 0.88534 * _A_BOHR / (zs[ii] ** 0.23 + zs[jj] ** 0.23) + phi = sum( + a_k * math.exp(-b_k * (r / a_screen)) + for a_k, b_k in zip(_ZBL_A_COEFF, _ZBL_B_COEFF, strict=True) + ) + total += _KE_EV_A * zs[ii] * zs[jj] / r * phi + return total + + +def _build_model_dict() -> dict: + """Build the combined native-spin + bridged dpmodel and jitter it. + + Returns + ------- + dict + The jittered serialized model tree, ready to be frozen. + """ + from deepmd.dpmodel.model.model import ( + get_model, + ) + + model = get_model(copy.deepcopy(SPIN_ZBL_CONFIG)) + assert model.has_spin() is True + kinds = [type(child).__name__ for child in model.atomic_model.models] + assert kinds[1] == "InterPotentialAtomicModel", ( + f"expected the bridged composition's second child to be the " + f"analytical InterPotentialAtomicModel, got {kinds!r}; without it " + f"this fixture is just the plain native-spin model again." + ) + model_dict = model.serialize() + model_dict = jitter_zero_arrays(model_dict, np.random.default_rng(_JITTER_SEED)) + return model_dict + + +def _assert_zbl_term_is_active(model_dict: dict) -> float: + """Pin that the analytical term genuinely contributes, and by how much. + + Re-builds the jittered model and its LEARNED child alone (the same model + minus the analytical energy), and checks their gas-phase energy difference + against the independent double-loop ZBL reference above. + + The identity is exact rather than approximate, but it is not simply + ``delta == ZBL``: ``jitter_zero_arrays`` also perturbs the all-zero + ``out_bias`` of the analytical child and of the linear composition itself, + and neither of those two per-type constants exists in the learned-child- + only model. Both are read straight out of the serialized tree and added + to the reference, so the check stays a bit-level identity (1e-10) instead + of a hand-tuned tolerance. + + Gas phase (no box) is used because the double-loop reference has no + periodic images; with this fixture's 6 A cell and 4 A cutoff the PBC case + would additionally involve images and a binding ``sel``. + + Parameters + ---------- + model_dict : dict + The jittered serialized model tree. + + Returns + ------- + float + The EAGER gas-phase total energy of the bridged model, so the caller + can hold the frozen archive to it (see ``main``). + """ + from deepmd.dpmodel.model.base_model import ( + BaseModel, + ) + from deepmd.dpmodel.model.native_spin_model import ( + NativeSpinEnergyModel, + ) + + model = BaseModel.deserialize(copy.deepcopy(model_dict)) + # Same wrapper class, same learned child, no analytical term: the ONLY + # difference from ``model`` is the composition's second child. + learned_only = NativeSpinEnergyModel( + atomic_model_=model.atomic_model.models[0], spin=model.spin + ) + # The dpmodel ``call`` API is framed (nf x nloc x ...); DeepPot.eval below + # takes the flat atype instead, hence the reshape here only. + atype_framed = _ATYPES.reshape(1, _NATOMS) + e_bridged = float( + np.reshape(model.call(_COORDS, atype_framed, _SPINS)["energy"], (-1,))[0] + ) + e_learned = float( + np.reshape(learned_only.call(_COORDS, atype_framed, _SPINS)["energy"], (-1,))[0] + ) + delta = e_bridged - e_learned + + zbl_ref = _analytic_zbl_total( + _COORDS[0], _ATYPES, float(SPIN_ZBL_CONFIG["descriptor"]["rcut"]) + ) + # The two per-type constant offsets that live only in the composition. + bias_inter = np.reshape(np.asarray(model.atomic_model.models[1].out_bias), (-1,)) + bias_linear = np.reshape(np.asarray(model.atomic_model.out_bias), (-1,)) + offset = float(np.sum(bias_inter[_ATYPES]) + np.sum(bias_linear[_ATYPES])) + + print( # noqa: T201 + f"\n// gas-phase bridged energy: {e_bridged:.18e}\n" + f"// gas-phase learned energy: {e_learned:.18e}\n" + f"// delta: {delta:.18e}\n" + f"// analytic ZBL + jittered out_bias offset: " + f"{zbl_ref + offset:.18e}" + ) + # (a) the analytical term must be LARGE -- it is the whole point of the + # fixture that a bridged energy is nowhere near the unbridged one. + assert abs(delta) > 1e-3, ( + f"bridging left the energy essentially unchanged (delta = {delta:.3e}); " + f"the analytical ZBL term is not reaching the composition, so this " + f"fixture would be vacuous." + ) + # (b) and it must be exactly the ZBL energy (plus the two jitter-induced + # constants), not merely "something large". + np.testing.assert_allclose( + delta, + zbl_ref + offset, + rtol=1e-10, + atol=1e-10, + err_msg=( + "the bridged-minus-learned energy is not the analytical ZBL sum; " + "the composition's analytical child is computing something else." + ), + ) + # (c) the learned half must not be swamped into irrelevance either: a + # fixture whose learned energy were zero would not test the DPA4 half. + assert abs(e_learned) > 1e-6, ( + f"learned-child energy is {e_learned:.3e}; the jitter did not take " + f"effect and the fixture would only exercise the analytical term." + ) + return e_bridged + + +def _check_metadata(pt2_path: str) -> None: + """Assert the frozen archive's metadata and the with-comm ABSENCE. + + Parameters + ---------- + pt2_path : str + Path to the frozen ``.pt2`` archive. + """ + with zipfile.ZipFile(pt2_path) as zf: + md = json.loads(zf.read("model/extra/metadata.json").decode("utf-8")) + names = zf.namelist() + print("\n// metadata:") # noqa: T201 + print( # noqa: T201 + json.dumps( + { + k: md[k] + for k in ( + "type_map", + "lower_input_kind", + "is_spin", + "has_comm_artifact", + "has_message_passing", + "use_spin", + "output_keys", + ) + if k in md + }, + indent=2, + ) + ) + assert md["type_map"] == SPIN_ZBL_CONFIG["type_map"] + assert md["lower_input_kind"] == "graph", ( + f"expected native-spin DPA4 to freeze to the graph lower, got " + f"{md.get('lower_input_kind')!r}" + ) + # Without is_spin the C++ side would route through DeepPot and the whole + # DeepSpin regression that consumes this archive would test nothing. + assert md["is_spin"] is True, ( + f"{pt2_path}: metadata is_spin = {md.get('is_spin')!r}, expected True; " + f"the composition dropped the spin flag on the way to the freeze." + ) + assert md["use_spin"] == [True, False] + # Single-rank only -- see the module docstring (the bridging gate's eta is + # incomplete per rank). BOTH halves matter: the flag is what the C++ + # dispatch reads, the archive entry is what it would load. + assert md["has_comm_artifact"] is False, ( + f"{pt2_path}: metadata has_comm_artifact = " + f"{md.get('has_comm_artifact')!r}, expected False; a bridged model " + f"cannot support multi-rank message passing (its Source Freeze " + f"Propagation Gate folds each node's full outgoing-edge set, which no " + f"rank owns), so advertising one would promise a capability the model " + f"cannot honour." + ) + assert "model/extra/forward_lower_with_comm.pt2" not in names, ( + f"{pt2_path}: a nested forward_lower_with_comm.pt2 was exported for a " + f"bridged model; see above -- the archive must not carry one." + ) + # The descriptor still message-passes WITHIN a rank; it is only the + # cross-rank exchange that bridging forbids. This is the flag that makes + # the C++ side fail fast under mpirun instead of answering wrongly. + assert md["has_message_passing"] is True + for key in ("atom_energy", "energy", "force", "force_mag", "virial"): + assert key in md["output_keys"] + + +def _eval_and_check(dp, cell, label: str) -> dict: + """Evaluate one (PBC or NoPbc) case and run its anti-vacuity checks. + + Parameters + ---------- + dp : deepmd.infer.DeepPot + The loaded archive. + cell : np.ndarray or None + The simulation cell, or ``None`` for a gas-phase evaluation. + label : str + Human-readable case name, used in messages. + + Returns + ------- + dict + The reference arrays for this case, in ``write_expected_ref``'s + field convention. + """ + e, f, v, ae, av, fm, _mm = dp.eval(_COORDS, cell, _ATYPES, atomic=True, spin=_SPINS) + print(f"\n// {label} total energy: {e[0, 0]:.18e}") # noqa: T201 + + assert np.all(np.isfinite(e)), f"{label}: non-finite energy" + assert np.all(np.isfinite(f)), f"{label}: non-finite force" + assert np.all(np.isfinite(fm)), f"{label}: non-finite force_mag" + + fmax = float(np.max(np.abs(f))) + print(f"// max |force|: {fmax:.6e}") # noqa: T201 + # Anti-vacuity: the 0.9 A Ni-Ni pair drives a large analytical ZBL + # repulsion, so a small max-force means the term (or the jitter) is not + # reaching the forward. + assert fmax > 1e-3, ( + f"{label}: expected a non-trivial force (the 0.9 A Ni-Ni pair drives " + f"the analytical ZBL term); got {fmax:.3e}." + ) + + spin_mask = _ATYPES == 0 # Ni carries spin; O does not + fm_flat = fm.reshape(_NATOMS, 3) + fm_spin_max = float(np.max(np.abs(fm_flat[spin_mask]))) + fm_nospin_max = float(np.max(np.abs(fm_flat[~spin_mask]))) + print( # noqa: T201 + f"// max |force_mag| spin / non-spin atoms: " + f"{fm_spin_max:.6e} / {fm_nospin_max:.6e}" + ) + # Anti-vacuity: a fresh (non-jittered) DPA4 zero-initializes its residual + # projections, making force_mag identically zero on the spin-carrying + # atoms too (see the jitter docstring in the module header). + assert fm_spin_max > 1e-6, ( + f"{label}: expected non-trivial force_mag on spin-active (Ni) atoms; " + f"got max |force_mag| = {fm_spin_max:.3e} (jitter not effective -- " + f"this fixture would be vacuous)." + ) + # The non-spin (O) rows must be EXACTLY gated to zero by the model's own + # type mask -- not merely small. The analytical child, which knows + # nothing about spin, must not leak into this channel either. + assert fm_nospin_max == 0.0, ( + f"{label}: expected force_mag to be EXACTLY zero on non-spin (O) " + f"atoms; got max |force_mag| = {fm_nospin_max:.3e}." + ) + return { + "expected_e": ae[0, :, 0], + "expected_f": f[0], + "expected_fm": fm[0], + "expected_tot_v": v[0], + "expected_atom_v": av[0], + } + + +def main(): + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file as pt_expt_deserialize_to_file, + ) + + ensure_inductor_compiler() + load_custom_ops() + + base_dir = os.path.dirname(__file__) + pt2_path = os.path.join(base_dir, "deeppot_dpa4_spin_zbl_graph.pt2") + ref_path = os.path.join(base_dir, "deeppot_dpa4_spin_zbl_graph.expected") + + # ---- 1. Build the jittered dpmodel dict from config+seed ---- + model_dict = _build_model_dict() + + # ---- 2. Pin that the analytical term is genuinely active (eager) ---- + # Done BEFORE the (slow) inductor compile so a degenerate fixture fails in + # seconds rather than minutes. + e_gas_eager = _assert_zbl_term_is_active(model_dict) + + data = { + "model": model_dict, + "model_def_script": SPIN_ZBL_CONFIG, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + + # ---- 3. Freeze directly to graph-kind .pt2 ---- + # Native-spin DPA4 has NO dense/nlist lower at all (spin rides the + # NeighborGraph lower exclusively), and the analytical child is likewise + # graph-route only, so ``lower_kind="auto"`` would resolve to "graph" + # anyway; pinned explicitly here for clarity. + print(f"\nExporting to {pt2_path} (lower_kind='graph') ...") # noqa: T201 + pt_expt_deserialize_to_file( + pt2_path, data, do_atomic_virial=True, lower_kind="graph" + ) + print("Export done.") # noqa: T201 + _check_metadata(pt2_path) + + # ---- 4. Evaluate the two reference cases ---- + from deepmd.infer import ( + DeepPot, + ) + + dp = DeepPot(pt2_path) + assert dp.has_spin + pbc = _eval_and_check(dp, _CELL, "PBC") + nopbc = _eval_and_check(dp, None, "NoPbc") + + # ---- 5. The freeze must preserve the composition ---- + # The checks above only prove the FROZEN archive is self-consistent; this + # holds it to the eager dpmodel it was built from, so a freeze that + # silently dropped the analytical child (or the spin injection) is caught + # here rather than becoming the "reference" every downstream C++/LAMMPS + # test then agrees with. Gas phase, to compare against the same eager + # evaluation the ZBL identity above used. 1e-10 is the project's + # cross-backend fp64 bound (dpmodel/NumPy vs the compiled torch artifact); + # the observed gap is ~1e-13 relative. + e_gas_frozen = float(np.sum(nopbc["expected_e"])) + np.testing.assert_allclose( + e_gas_frozen, + e_gas_eager, + rtol=1e-10, + atol=0.0, + err_msg=( + "the frozen archive's gas-phase energy disagrees with the eager " + "dpmodel it was built from; the export dropped part of the " + "native-spin + bridging composition." + ), + ) + + # ---- 6. Sidecar reference consumed by the C++ test ---- + write_expected_ref( + ref_path, + sections={"pbc": pbc, "nopbc": nopbc}, + source_script="source/tests/infer/gen_dpa4_spin_zbl.py", + ) + print(f"\nWrote {ref_path}") # noqa: T201 + + print("\nDone!") # noqa: T201 + + +if __name__ == "__main__": + main() diff --git a/source/tests/infer/gen_dpa4_zbl.py b/source/tests/infer/gen_dpa4_zbl.py new file mode 100644 index 0000000000..a8f5f6ade0 --- /dev/null +++ b/source/tests/infer/gen_dpa4_zbl.py @@ -0,0 +1,204 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Generate deeppot_dpa4_zbl_graph.pt2: DPA4 with analytical ZBL bridging. + +``bridging_method: ZBL`` builds a COMPOSITION -- ``LinearEnergyModel`` over +``[learned DPA4, InterPotentialAtomicModel]`` with ``weights="sum"`` -- so +the frozen archive exercises a code path no other C++ fixture covers: the +graph lower of a linear composition rather than a single learned model. +Before this fixture, ZBL bridging had NO C++ or LAMMPS coverage at all; its +only end-to-end test drove the ``.pt2`` through the PYTHON ``DeepPot``, which +never touches ``DeepPotPTExpt``. + +Single-rank only: bridging enables the descriptor's Source Freeze +Propagation Gate, whose per-node ``eta_j = prod_{e: src_e = j} w_e`` folds a +node's FULL outgoing-edge set. Edges exist only for owned centres, so eta is +incomplete on every rank and a with-comm artifact is not exported +(``has_comm_artifact=false``, asserted below). + +Generation mirrors ``gen_dpa4_spin.py``: the dpmodel is built in-process from +the inline config with a fixed weight-init seed, its zero-initialised +residual projections are jittered away from exact zero with a fixed RNG seed, +and the result is frozen straight to the graph-kind ``.pt2``. Without the +jitter a fresh DPA4 collapses to a type-embedding-only descriptor and every +force would be identically zero, making the fixture vacuous. +""" + +import copy +import json +import os +import sys +import zipfile + +import numpy as np + +# Ensure the source tree is on the path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..")) +# Ensure source/tests is on the path for dpa4_fixtures (this script runs +# standalone, outside pytest's package machinery). +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from dpa4_fixtures import ( + jitter_zero_arrays, +) +from gen_common import ( + ensure_inductor_compiler, + load_custom_ops, + write_expected_ref, +) + +# Small fp64 DPA4 + ZBL config. ``bridging_r_inner``/``r_outer`` feed the +# descriptor's InnerClamp AND BridgingSwitch (they are built together from +# the same radii), and the model-level InterPotential term. +ZBL_CONFIG = { + "type_map": ["Ni", "O"], + "descriptor": { + "type": "dpa4", + "rcut": 4.0, + "sel": 8, + "channels": 16, + "n_radial": 8, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "precision": "float64", + "seed": 7, + "random_gamma": False, + }, + "fitting_net": { + "type": "dpa4_ener", + "neuron": [8, 8], + "precision": "float64", + "seed": 7, + }, + "bridging_method": "ZBL", + "bridging_r_inner": 0.8, + "bridging_r_outer": 1.2, +} + +_JITTER_SEED = 20260725 + +# Fixed 6-atom system. Atoms 0 and 1 sit 0.9 A apart -- inside +# ``bridging_r_outer`` -- so the analytical ZBL term contributes a large, +# unmistakable repulsion instead of a numerical afterthought. +_NATOMS = 6 +_ATYPES = np.array([0, 0, 0, 1, 1, 1], dtype=np.int32) # Ni, Ni, Ni, O, O, O +_COORDS = np.array( + [ + [1.0, 1.0, 1.0], + [1.9, 1.0, 1.0], + [1.3, 1.8, 1.0], + [0.4, 1.2, 1.6], + [3.6, 2.0, 1.3], + [3.4, 0.7, 1.7], + ], + dtype=np.float64, +).reshape(1, _NATOMS, 3) +_CELL = (np.eye(3, dtype=np.float64) * 6.0).reshape(1, 9) + + +def main(): + from deepmd.dpmodel.model.model import ( + get_model, + ) + from deepmd.infer import ( + DeepPot, + ) + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file as pt_expt_deserialize_to_file, + ) + + ensure_inductor_compiler() + load_custom_ops() + + base_dir = os.path.dirname(__file__) + pt2_path = os.path.join(base_dir, "deeppot_dpa4_zbl_graph.pt2") + + # ---- 1. Build the jittered composition ---- + model = get_model(copy.deepcopy(ZBL_CONFIG)) + model_dict = jitter_zero_arrays( + model.serialize(), np.random.default_rng(_JITTER_SEED) + ) + data = { + "model": model_dict, + "model_def_script": ZBL_CONFIG, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + + # ---- 2. Freeze to graph-kind .pt2 ---- + print(f"Exporting to {pt2_path} (lower_kind='graph') ...") # noqa: T201 + pt_expt_deserialize_to_file( + pt2_path, data, do_atomic_virial=True, lower_kind="graph" + ) + print("Export done.") # noqa: T201 + + # ---- 3. Check metadata ---- + with zipfile.ZipFile(pt2_path) as zf: + md = json.loads(zf.read("model/extra/metadata.json").decode("utf-8")) + names = zf.namelist() + print( # noqa: T201 + json.dumps( + { + k: md[k] + for k in ("type_map", "lower_input_kind", "has_comm_artifact") + if k in md + }, + indent=2, + ) + ) + assert md["type_map"] == ZBL_CONFIG["type_map"] + assert md["lower_input_kind"] == "graph" + # Single-rank only -- see the module docstring (SFPG eta is incomplete + # per rank), so no nested with-comm artifact may be present. + assert md["has_comm_artifact"] is False + assert "model/extra/forward_lower_with_comm.pt2" not in names + + # ---- 4. Evaluate (PBC + NoPbc) ---- + dp = DeepPot(pt2_path) + e1, f1, v1, ae1, av1 = dp.eval(_COORDS, _CELL, _ATYPES, atomic=True) + e_np, f_np, v_np, ae_np, av_np = dp.eval(_COORDS, None, _ATYPES, atomic=True) + print(f"\n// PBC total energy: {e1[0, 0]:.18e}") # noqa: T201 + print(f"// NoPbc total energy: {e_np[0, 0]:.18e}") # noqa: T201 + + for label, e, f in (("PBC", e1, f1), ("NoPbc", e_np, f_np)): + assert np.all(np.isfinite(e)), f"{label}: non-finite energy" + assert np.all(np.isfinite(f)), f"{label}: non-finite force" + fmax = float(np.max(np.abs(f))) + print(f"// {label} max |force|: {fmax:.6e}") # noqa: T201 + # Anti-vacuity: an unjittered DPA4 gives identically zero forces. + # The close Ni-Ni pair alone drives a large ZBL repulsion, so a + # small max-force means the fixture is degenerate. + assert fmax > 1e-3, ( + f"{label}: expected a non-trivial force (the 0.9 A Ni-Ni pair " + f"drives the analytical ZBL term); got {fmax:.3e} -- jitter or " + f"bridging is not effective and this fixture would be vacuous." + ) + + # ---- 5. Sidecar reference consumed by the C++ test ---- + ref_path = os.path.join(base_dir, "deeppot_dpa4_zbl_graph.expected") + write_expected_ref( + ref_path, + sections={ + "pbc": { + "expected_e": ae1[0, :, 0], + "expected_f": f1[0], + "expected_tot_v": v1[0], + "expected_atom_v": av1[0], + }, + "nopbc": { + "expected_e": ae_np[0, :, 0], + "expected_f": f_np[0], + "expected_tot_v": v_np[0], + "expected_atom_v": av_np[0], + }, + }, + source_script="source/tests/infer/gen_dpa4_zbl.py", + ) + print(f"Wrote {ref_path}") # noqa: T201 + print("\nDone!") # noqa: T201 + + +if __name__ == "__main__": + main() diff --git a/source/tests/pt/model/test_dpa4_dpmodel_parity.py b/source/tests/pt/model/test_dpa4_dpmodel_parity.py index feeba7d56d..dab9fd6833 100644 --- a/source/tests/pt/model/test_dpa4_dpmodel_parity.py +++ b/source/tests/pt/model/test_dpa4_dpmodel_parity.py @@ -3010,6 +3010,76 @@ def _build_real_edge_inputs( } +def _dp_cache_from_padded( + *, + type_ebed, + coord, + nlist, + mapping, + pair_keep_mask, + eps, + deg_norm_floor, + edge_envelope, + radial_basis, + random_gamma, + wigner_calc, + gamma=None, +): + """Build the dp ``EdgeCache`` from a padded quartet via the surviving seam. + + The retired padded dense builder (``build_edge_cache``) is gone; the + padded-quartet contract it pinned now lives in + ``DescrptDPA4._graph_from_padded_nlist`` (``src_ok`` sanitization) + + ``graph_from_dense_quartet`` (row-major shape-static conversion) + + ``_edge_cache_from_arrays`` (the one edge-native cache core). This + helper chains them exactly as the production dense adapter does, so the + padded parity tests keep pinning the same contract against the surviving + architecture. + + The fixture's ``pair_keep_mask`` (which the dense builder folded into its + validity mask) is ANDed into the graph's ``edge_mask`` before the core -- + mirroring what the canonical ``apply_pair_exclusion`` transform does, + matching the production graph route's single-exclusion-site contract + (the core itself has no exclusion parameter and never re-applies it). + ``_edge_cache_from_arrays`` folds masking into the per-edge weights and + leaves ``edge_mask`` unset; the padded tests assert on the validity + mask, so it is attached to the returned cache. + """ + import dataclasses + + from deepmd.dpmodel.descriptor.dpa4 import ( + _graph_from_padded_nlist, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import ( + _edge_cache_from_arrays, + ) + + nf, nloc, _ = nlist.shape + nall = coord.shape[1] + # atype_ext only feeds the converter's (unused here) flat-local-atype + # output; pair exclusion enters via the fixture's pair_keep_mask below. + atype_ext_dummy = np.zeros((nf, nall), dtype=np.int64) + graph, _ = _graph_from_padded_nlist(coord, atype_ext_dummy, nlist, mapping) + edge_mask = np.asarray(graph.edge_mask) & pair_keep_mask.reshape(-1) + cache = _edge_cache_from_arrays( + type_ebed=type_ebed, + edge_index=graph.edge_index, + edge_vec=graph.edge_vec, + edge_mask=edge_mask, + compute_dtype=np.float64, + eps=eps, + deg_norm_floor=deg_norm_floor, + inner_clamp=None, + bridging_switch=None, + edge_envelope=edge_envelope, + radial_basis=radial_basis, + random_gamma=random_gamma, + wigner_calc=wigner_calc, + gamma=gamma, + ) + return dataclasses.replace(cache, edge_mask=edge_mask) + + def _build_real_edge_caches( inputs, *, @@ -3022,15 +3092,12 @@ def _build_real_edge_caches( gamma=None, seed=2090, ): - """Run the REAL pt and dp ``build_edge_cache`` on identical inputs. + """Run the REAL pt builder and the surviving dp seam on identical inputs. The pt ``RadialBasis`` frequencies are perturbed and weight-copied into the dp side via ``deserialize`` so parity exercises copied weights. Returns ``(pt_cache, dp_cache)``. """ - from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import ( - build_edge_cache as dp_build_edge_cache, - ) from deepmd.dpmodel.descriptor.dpa4_nn.radial import C3CutoffEnvelope as DPEnvelope from deepmd.dpmodel.descriptor.dpa4_nn.radial import RadialBasis as DPRadialBasis from deepmd.dpmodel.descriptor.dpa4_nn.wignerd import WignerDCalculator as DPWigner @@ -3067,9 +3134,9 @@ def _build_real_edge_caches( random_gamma=random_gamma, wigner_calc=pt_wig, ) - dp_cache = dp_build_edge_cache( + dp_cache = _dp_cache_from_padded( type_ebed=inputs["type_ebed"], - extended_coord=inputs["coord"], + coord=inputs["coord"], nlist=inputs["nlist"], mapping=mapping, pair_keep_mask=inputs["pair_keep_mask"], @@ -3077,7 +3144,6 @@ def _build_real_edge_caches( deg_norm_floor=deg_norm_floor, edge_envelope=dp_env, radial_basis=dp_rb, - n_radial=n_radial, random_gamma=random_gamma, wigner_calc=dp_wig, gamma=gamma, @@ -3170,11 +3236,8 @@ def test_real_build_parity(self, local_nlist, deg_norm_floor) -> None: def test_out_of_range_local_index_masked(self) -> None: # a local nlist entry >= nloc with mapping=None must be masked out - # and must not break the coordinate gather (nlist_safe is re-zeroed - # after the final src_ok mask update) - from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import ( - build_edge_cache as dp_build_edge_cache, - ) + # and must not break the coordinate gather (the src_ok sanitization + # rewrites the slot to -1 before conversion) from deepmd.dpmodel.descriptor.dpa4_nn.radial import ( C3CutoffEnvelope as DPEnvelope, ) @@ -3189,9 +3252,9 @@ def test_out_of_range_local_index_masked(self) -> None: nlist = inputs["nlist"].copy() nlist[0, 0, 0] = self.nloc # out of [0, nloc), would gather OOB n_radial = 8 - cache = dp_build_edge_cache( + cache = _dp_cache_from_padded( type_ebed=inputs["type_ebed"], - extended_coord=inputs["coord"], + coord=inputs["coord"], nlist=nlist, mapping=None, pair_keep_mask=inputs["pair_keep_mask"], @@ -3201,7 +3264,6 @@ def test_out_of_range_local_index_masked(self) -> None: radial_basis=DPRadialBasis( rcut=6.0, n_radial=n_radial, precision="float64" ), - n_radial=n_radial, random_gamma=False, wigner_calc=DPWigner(self.lmax, precision="float64"), ) diff --git a/source/tests/pt_expt/descriptor/test_dpa4.py b/source/tests/pt_expt/descriptor/test_dpa4.py index ef4a10227b..ca583ebe6a 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4.py +++ b/source/tests/pt_expt/descriptor/test_dpa4.py @@ -1,5 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from unittest import ( + mock, +) + import numpy as np import pytest import torch @@ -99,6 +103,101 @@ def test_consistency(self, use_env_seed, use_mapping) -> None: err_msg=err_msg, ) + def test_random_gamma_train_eval_gate(self) -> None: + """``random_gamma`` mirrors pt: rolled in train mode, fixed otherwise. + + The ``_in_training_mode`` runtime hook must forward + ``random_gamma=True`` to the edge-cache builder only for a + train-mode pt_expt forward; eval-mode forwards and the dpmodel + reference (no training mode) always forward ``False``. Spy-based: + the model is roll-equivariant, so an output-difference check would + have no deterministic teeth. + """ + import deepmd.dpmodel.descriptor.dpa4 as dpa4_mod + + dtype = PRECISION_DICT["float64"] + dd0 = make_descriptor( + self.nt, + self.sel_mix, + self.rcut, + random_gamma=True, # the default in production configs + ).to(self.device) + coord_ext = torch.tensor(self.coord_ext, dtype=dtype, device=self.device) + atype_ext = torch.tensor(self.atype_ext, dtype=int, device=self.device) + nlist = torch.tensor(self.nlist, dtype=int, device=self.device) + + captured: list[bool] = [] + orig = dpa4_mod._edge_cache_from_arrays + + def spy(*args, **kwargs): + captured.append(kwargs["random_gamma"]) + return orig(*args, **kwargs) + + with mock.patch.object(dpa4_mod, "_edge_cache_from_arrays", spy): + # Train mode: the augmentation is on (and the numpy gamma draw + # runs against torch tensors -- eager smoke). + dd0.train() + dd0(coord_ext, atype_ext, nlist) + assert captured[-1] is True + # Eval mode: fixed gamma, deterministic forwards. + dd0.eval() + r1 = dd0(coord_ext, atype_ext, nlist)[0] + r2 = dd0(coord_ext, atype_ext, nlist)[0] + assert captured[-1] is False + assert torch.equal(r1, r2) + # dpmodel reference: no training mode, never rolls even with + # random_gamma=True. + dd2 = DPDescrptDPA4.deserialize(dd0.serialize()) + assert dd2._in_training_mode() is False + dd2.call(self.coord_ext, self.atype_ext, self.nlist) + assert captured[-1] is False + + def test_random_gamma_draw_obeys_torch_rng_state(self) -> None: + """The train-mode roll must be drawn by torch, on the edge device. + + pt draws it with ``torch.rand``, so ``torch.manual_seed`` replays it. + A host numpy draw would answer to no seed and freeze to a constant + under tracing. Replay is the only property separating the two: the + descriptor is roll-equivariant, so comparing outputs has no teeth. + """ + import deepmd.dpmodel.descriptor.dpa4_nn.edge_cache as ec_mod + + dtype = PRECISION_DICT["float64"] + dd0 = make_descriptor(self.nt, self.sel_mix, self.rcut, random_gamma=True).to( + self.device + ) + dd0.train() + coord_ext = torch.tensor(self.coord_ext, dtype=dtype, device=self.device) + atype_ext = torch.tensor(self.atype_ext, dtype=int, device=self.device) + nlist = torch.tensor(self.nlist, dtype=int, device=self.device) + + def _rolled_quat() -> torch.Tensor: + """One train-mode forward's post-roll edge quaternion.""" + captured: list[torch.Tensor] = [] + orig = ec_mod.quaternion_multiply + + def spy(a, b): + out = orig(a, b) + captured.append(out) + return out + + with mock.patch.object(ec_mod, "quaternion_multiply", spy): + dd0(coord_ext, atype_ext, nlist) + assert captured, "the random roll never ran" + return captured[0].detach().clone() + + torch.manual_seed(20260725) + q_a = _rolled_quat() + torch.manual_seed(20260725) + q_b = _rolled_quat() + torch.testing.assert_close(q_a, q_b, rtol=0.0, atol=0.0) + + # anti-vacuity: a constant gamma would satisfy the replay check alone + torch.manual_seed(20260726) + assert not torch.allclose(q_a, _rolled_quat()) + # drawn on the edge device, no host round-trip + assert q_a.device.type == self.device.type + @pytest.mark.parametrize("prec", ["float64"]) # precision def test_exportable(self, prec) -> None: dtype = PRECISION_DICT[prec] diff --git a/source/tests/pt_expt/infer/test_dpa4_deep_eval.py b/source/tests/pt_expt/infer/test_dpa4_deep_eval.py index bcade0d77e..eb9401c826 100644 --- a/source/tests/pt_expt/infer/test_dpa4_deep_eval.py +++ b/source/tests/pt_expt/infer/test_dpa4_deep_eval.py @@ -192,3 +192,54 @@ def test_dpa4_deep_eval_metadata(dpa4_pt_and_pt2) -> None: assert dp_pt2.deep_eval.get_dim_fparam() == dp_pt.deep_eval.get_dim_fparam() assert dp_pt2.deep_eval.get_dim_aparam() == dp_pt.deep_eval.get_dim_aparam() assert not dp_pt2.has_spin + + +# --------------------------------------------------------------------------- +# Graph-route parity: persisted fixtures from source/tests/infer/gen_dpa4.py +# (Section B), independent of the pt-vs-pt_expt fixture above. +# --------------------------------------------------------------------------- + +_INFER_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "infer") +_DPA4_GRAPH_PT2 = os.path.join(_INFER_DIR, "deeppot_dpa4_graph.pt2") +_DPA4_GRAPH_NLIST_REF_PT2 = os.path.join(_INFER_DIR, "deeppot_dpa4_graph_nlist_ref.pt2") + + +@pytest.mark.skipif( + not (os.path.exists(_DPA4_GRAPH_PT2) and os.path.exists(_DPA4_GRAPH_NLIST_REF_PT2)), + reason="gen_dpa4.py graph fixtures not generated (e.g. skipped under " + "LeakSanitizer, or gen_dpa4.py has not been run)", +) +@pytest.mark.parametrize("pbc", [True, False]) # periodic vs open boundary +def test_dpa4_graph_deep_eval_matches_nlist_ref(pbc) -> None: + """DPA4 graph ``.pt2`` (persisted by gen_dpa4.py Section B) vs its + independent dense-nlist oracle (same jittered weights, + ``lower_kind="nlist"``), both reloaded through :class:`DeepPot`. + + gen_dpa4.py already performs this comparison in-process at generation + time (``cross_tol=1e-8``, mirroring gen_dpa2.py's B.4); this test + instead exercises the persisted-artifact reload path through the public + ``DeepPot`` API -- a regression test for the on-disk ``.pt2`` format, + independent of the in-process objects used during generation. Both + artifacts are fp64 end-to-end (descriptor/fitting ``precision: + "float64"``), so cross-path energy/force agreement is expected at the + same noise floor as the gen-time check; the same ``rtol=atol=1e-8`` + threshold is reused here (not invented) for consistency. The per-atom + virial is NOT compared: the graph path assigns each edge's force/virial + contribution fully to the source atom (``edge_force_virial`` + full-to-src), a different (equally valid) decomposition than the dense + per-atom one -- only the sum (global virial, which IS compared here) is + convention-independent. + """ + dp_graph = DeepPot(_DPA4_GRAPH_PT2) + dp_nlist = DeepPot(_DPA4_GRAPH_NLIST_REF_PT2) + + cell = _CELL if pbc else None + e_g, f_g, v_g, ae_g, av_g = dp_graph.eval(_COORDS, cell, _ATYPES, atomic=True) + e_n, f_n, v_n, ae_n, av_n = dp_nlist.eval(_COORDS, cell, _ATYPES, atomic=True) + + tag = "pbc" if pbc else "nopbc" + cross_tol = {"rtol": 1e-8, "atol": 1e-8} + np.testing.assert_allclose(e_g, e_n, err_msg=f"{tag}: energy", **cross_tol) + np.testing.assert_allclose(f_g, f_n, err_msg=f"{tag}: force", **cross_tol) + np.testing.assert_allclose(v_g, v_n, err_msg=f"{tag}: global virial", **cross_tol) + np.testing.assert_allclose(ae_g, ae_n, err_msg=f"{tag}: atom energy", **cross_tol) diff --git a/source/tests/pt_expt/infer/test_edge_vec_deprecation.py b/source/tests/pt_expt/infer/test_edge_vec_deprecation.py new file mode 100644 index 0000000000..90ac4a8a92 --- /dev/null +++ b/source/tests/pt_expt/infer/test_edge_vec_deprecation.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""The edge_vec .pt2 schema (pt-backend SeZM freeze) is a deprecated +legacy format; loading one must warn, loading anything else must not. +""" + +import warnings + +import pytest + + +def test_edge_vec_metadata_warns() -> None: + from deepmd.pt_expt.infer.deep_eval import ( + _warn_legacy_edge_vec, + ) + + with pytest.warns(DeprecationWarning, match="edge_vec"): + _warn_legacy_edge_vec({"lower_input_kind": "edge_vec"}) + + +@pytest.mark.parametrize( + "kind", + [ + "nlist", # pt_expt dense freeze + "graph", # pt_expt graph freeze + "dpa1_canonical", # dpa1 compact-canonical freeze + None, # pre-metadata archives + ], +) +def test_other_kinds_do_not_warn(kind) -> None: + from deepmd.pt_expt.infer.deep_eval import ( + _warn_legacy_edge_vec, + ) + + with warnings.catch_warnings(): + warnings.simplefilter("error") + _warn_legacy_edge_vec({"lower_input_kind": kind} if kind else {}) diff --git a/source/tests/pt_expt/model/test_dpa4_export.py b/source/tests/pt_expt/model/test_dpa4_export.py index bb06b25574..dbe7c08b45 100644 --- a/source/tests/pt_expt/model/test_dpa4_export.py +++ b/source/tests/pt_expt/model/test_dpa4_export.py @@ -1,21 +1,33 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Model-level freeze test for the DPA4/SeZM energy model. -Mirrors the DPA3 ``test_export_with_comm`` round-trip: a DPA4 model is a -GNN (``has_message_passing_across_ranks() == True``), so -``deserialize_to_file`` produces a .pt2 archive containing TWO compiled -artifacts: - * the regular ``forward_lower`` (no comm), packed at the top of the ZIP; - * a ``forward_lower_with_comm`` variant nested at - ``model/extra/forward_lower_with_comm.pt2``. - -This test verifies: - 1. The .pt2 archive is produced and both artifacts are present. - 2. ``metadata.json`` carries the correct ``type_map``/``rcut`` and - ``has_message_passing: true`` (DPA4 is a message-passing descriptor). +A DPA4 model is a message-passing GNN (``has_message_passing() == True``), +and cross-rank ghost-feature exchange is implemented ONLY on the +NeighborGraph lower: it carries a real per-layer ``border_op`` exchange +(``has_message_passing_across_ranks()`` is True). The dense (nlist) lower's +``call`` adapter still raises on ``comm_dict`` (``dense_lower_supports_comm()`` +is False), so ``deserialize_to_file`` embeds the ``forward_lower_with_comm`` +sidecar for the ``"graph"`` kind only; the ``"nlist"`` kind stays a SINGLE +compiled artifact and multi-rank inference on it must fail fast at the C++ +dispatch instead of silently skipping the exchange. + +Task 8 parametrizes the freeze over ``lower_kind``: ``"auto"`` now resolves +to ``"graph"`` (``_resolve_lower_kind`` sees ``model_uses_graph_lower() is +True``; ``canonical_model_eligible`` is dpa1-specific so DPA4 never takes the +``"dpa1_canonical"`` branch), while ``"nlist"`` stays reachable for +back-compat. This test verifies, per ``lower_kind``: + 1. The .pt2 archive is produced and the with-comm artifact is PRESENT for + the ``"graph"`` kind (cross-rank exchange) and ABSENT for the + ``"nlist"`` kind (dense lower is comm-less). + 2. ``metadata.json`` carries the correct ``type_map``/``rcut``, + ``has_message_passing: true``, kind-conditional ``has_comm_artifact``, + and ``lower_input_kind == expected_input_kind``; the graph kind + additionally carries ``graph_edge_dtype``. 3. The regular artifact loads via ``aoti_load_package``. - 4. The loaded artifact's ``forward_common_lower`` output matches the - eager model (fp64 AOTI parity, rtol 1e-10). + 4. The loaded artifact reproduces the eager model: the ``"nlist"`` kind + against ``forward_common_lower`` (dense ABI, fp64 AOTI parity, rtol + 1e-10); the ``"graph"`` kind against ``forward_common_lower_graph`` + (NeighborGraph ABI, same tolerance). """ from __future__ import ( @@ -28,18 +40,49 @@ import numpy as np import pytest +import torch # Note: registration of the deepmd_export::border_op opaque wrapper (needed by # the with-comm artifact) happens inside ``deserialize_to_file`` via # ``ensure_comm_registered()``; no explicit comm import is required here. +from deepmd.pt_expt.model.ener_model import ( + _translate_energy_keys, +) from deepmd.pt_expt.model.get_model import ( get_model, ) +from deepmd.pt_expt.model.model import ( + BaseModel, +) +from deepmd.pt_expt.utils import env as _env from deepmd.pt_expt.utils.serialization import ( _make_sample_inputs, + build_synthetic_graph_inputs, deserialize_to_file, ) +from ...dpa4_fixtures import ( + jitter_zero_arrays, +) +from .test_dpa4_native_spin import ( + NATIVE_SPIN_CONFIG, + _build_native_spin_model_cpu, +) + + +def _to_artifact_device(*tensors: torch.Tensor | None) -> tuple: + """Move sample tensors to the AOTI artifact's compile device. + + ``deserialize_to_file`` runs ``move_to_device_pass(exported, _env.DEVICE)`` + before AOTI compile, so a CUDA box produces a CUDA-only artifact; feeding + it CPU tensors triggers an illegal-memory-access at the AOTI boundary. + The eager reference stays on CPU untouched -- only the artifact call + needs the move. ``None`` placeholders (unset fparam/aparam/charge_spin) + pass through unchanged. + """ + return tuple(t if t is None else t.to(_env.DEVICE) for t in tensors) + + # Small fp64 DPA4 config (channels 16, n_radial 8, lmax 2, mmax 1, # n_blocks 2) — large enough to exercise the SO(2)/SO(3) + attention + # embedding paths that previously specialized ``nloc`` during export, but @@ -72,17 +115,44 @@ os.environ.get("CI") == "true", reason="AOTInductor compile is slow (minutes); run locally only by default.", ) -def test_dpa4_freeze_to_pt2(tmp_path) -> None: - """End-to-end: DPA4 model freezes to a dual-artifact .pt2 and the - regular artifact reproduces the eager ``forward_common_lower``. +@pytest.mark.parametrize( + "lower_kind,expected_input_kind", + [ + ("auto", "graph"), # default now resolves to the graph lower + ("nlist", "nlist"), # dense kind stays reachable for back-compat + ], +) +def test_dpa4_freeze_to_pt2(tmp_path, lower_kind, expected_input_kind) -> None: + """End-to-end: DPA4 model freezes to a .pt2 archive with kind-conditional + artifact layout (with-comm sidecar embedded for the ``"graph"`` kind's + multi-rank exchange, single-artifact for the ``"nlist"`` kind's comm-less + dense lower), and the regular artifact reproduces the matching eager + forward (dense ``forward_common_lower`` for ``"nlist"``, + ``forward_common_lower_graph`` for ``"graph"``). """ model = get_model(_DPA4_CONFIG) model.to("cpu") model.eval() # 1. Serialize → deserialize_to_file (compiles and packs both artifacts). - pt2_path = str(tmp_path / "test_dpa4.pt2") - deserialize_to_file(pt2_path, {"model": model.serialize()}) + # + # DPA4 deliberately zero-initializes several residual output + # projections (see ``jitter_zero_arrays``'s docstring in + # ``dpa4_fixtures.py``), so a fresh, untrained model is architecturally + # edge-INDEPENDENT: force/virial are ~0 regardless of edge handling. + # That would make the AOTI-vs-eager parity below vacuous (it couldn't + # catch an inductor miscompile of the edge scatter/index_add path) for + # either lower kind, so the jitter is applied uniformly to both + # parametrizations. + data = {"model": model.serialize()} + data["model"] = jitter_zero_arrays(data["model"], np.random.default_rng(99)) + # Rebuild the eager reference model from the SAME jittered dict that + # gets frozen below, so the AOTI artifact and the eager reference + # compared against it are the same (edge-sensitive) model. + model = BaseModel.deserialize(data["model"]).to("cpu") + model.eval() + pt2_path = str(tmp_path / f"test_dpa4_{expected_input_kind}.pt2") + deserialize_to_file(pt2_path, data, lower_kind=lower_kind) assert os.path.exists(pt2_path) # 2. ZIP layout + metadata sanity. PyTorch's strict layout puts our @@ -90,14 +160,22 @@ def test_dpa4_freeze_to_pt2(tmp_path) -> None: with zipfile.ZipFile(pt2_path, "r") as zf: names = set(zf.namelist()) meta = json.loads(zf.read("model/extra/metadata.json").decode("utf-8")) - assert "model/extra/forward_lower_with_comm.pt2" in names, ( - f"with-comm artifact missing; names={sorted(names)}" - ) + if expected_input_kind == "graph": + assert "model/extra/forward_lower_with_comm.pt2" in names, ( + "graph kind must embed the with-comm sidecar (multi-rank)" + ) + assert meta["has_comm_artifact"] is True + else: + assert "model/extra/forward_lower_with_comm.pt2" not in names, ( + "nlist kind must stay single-artifact (dense lower is comm-less)" + ) + assert meta["has_comm_artifact"] is False assert meta["type_map"] == _DPA4_CONFIG["type_map"] assert meta["rcut"] == model.get_rcut() - # DPA4 is a message-passing GNN descriptor. + # DPA4 is a message-passing GNN descriptor; only the graph lower + # implements the cross-rank exchange (see module docstring). assert meta["has_message_passing"] is True - assert meta["has_comm_artifact"] is True + assert meta["lower_input_kind"] == expected_input_kind # 3. The regular artifact loads. from torch._inductor import ( @@ -106,43 +184,691 @@ def test_dpa4_freeze_to_pt2(tmp_path) -> None: regular = aoti_load_package(pt2_path) - # 4. Eager reference vs. AOTI artifact parity on forward_common_lower. - sample = _make_sample_inputs(model, nframes=1, has_spin=False) - ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin = sample - - eager_out = model.forward_common_lower( - ext_coord.detach().requires_grad_(True), - ext_atype, - nlist_t, - mapping_t, - fparam=fparam, - aparam=aparam, - do_atomic_virial=False, - charge_spin=charge_spin, - ) - - artifact_out = regular( - ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin - ) - - # The AOTI artifact returns the internal forward_common_lower keys; compare - # every key it produces against the eager reference (fp64 AOTI tolerance). - compared = 0 - for key, val in artifact_out.items(): - if key not in eager_out or eager_out[key] is None or val is None: - continue - np.testing.assert_allclose( - val.detach().cpu().numpy(), - eager_out[key].detach().cpu().numpy(), - rtol=1e-10, - atol=1e-10, - err_msg=f"artifact vs eager forward_common_lower differs: {key}", + if expected_input_kind == "nlist": + # 4a. Dense-ABI eager reference vs. AOTI artifact parity on + # forward_common_lower. + # _make_sample_inputs creates tensors on _env.DEVICE (CUDA on a GPU + # box); the eager reference model lives on CPU, so make explicit CPU + # copies for it. The artifact call below gets separate _env.DEVICE + # copies via _to_artifact_device. + sample = _make_sample_inputs(model, nframes=1, has_spin=False) + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin = tuple( + t if t is None else t.to("cpu") for t in sample + ) + + eager_out = model.forward_common_lower( + ext_coord.detach().requires_grad_(True), + ext_atype, + nlist_t, + mapping_t, + fparam=fparam, + aparam=aparam, + do_atomic_virial=False, + charge_spin=charge_spin, + ) + + # Anti-vacuity guard: fresh DPA4 is edge-independent (forces ~0); + # confirm the jitter above made the eager reference force + # non-trivial, else the AOTI parity below would pass trivially. + f_ref = eager_out["energy_derv_r"].detach().cpu().numpy() + assert np.abs(f_ref).max() > 1e-6, ( + f"eager reference force is near-zero ({np.abs(f_ref).max():.3e}); " + f"jitter not effective -- AOTI parity check would be vacuous" + ) + + # The artifact is compiled for _env.DEVICE (move_to_device_pass in + # deserialize_to_file); move inputs there while the eager reference + # above stays on CPU. + ( + d_ext_coord, + d_ext_atype, + d_nlist_t, + d_mapping_t, + d_fparam, + d_aparam, + d_charge_spin, + ) = _to_artifact_device( + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin + ) + artifact_out = regular( + d_ext_coord, + d_ext_atype, + d_nlist_t, + d_mapping_t, + d_fparam, + d_aparam, + d_charge_spin, + ) + + # The AOTI artifact returns the internal forward_common_lower keys; + # compare every key it produces against the eager reference (fp64 + # AOTI tolerance). + compared = 0 + for key, val in artifact_out.items(): + if key not in eager_out or eager_out[key] is None or val is None: + continue + np.testing.assert_allclose( + val.detach().cpu().numpy(), + eager_out[key].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + err_msg=f"artifact vs eager forward_common_lower differs: {key}", + ) + compared += 1 + # Guard against a vacuous pass (no overlapping keys compared). + assert compared > 0, ( + f"no overlapping output keys compared; artifact keys=" + f"{sorted(artifact_out)}, eager keys={sorted(eager_out)}" + ) + # The energy output must be among the compared keys. + assert "energy_redu" in artifact_out or "energy" in artifact_out + else: + # 4b. NeighborGraph-ABI eager reference vs. AOTI artifact parity on + # forward_common_lower_graph. Metadata carries the edge dtype the + # graph artifact was frozen with. + assert meta["graph_edge_dtype"] in ("float32", "float64") + + sample = build_synthetic_graph_inputs( + model, + e_max=None, + nframes=1, + nloc=6, + dtype=torch.float64, + device=torch.device("cpu"), + ) + atype, n_node, n_local, ei, ev, em, do, drp, so, srp, fp, ap, cs = sample + + eager_internal = model.forward_common_lower_graph( + atype, + n_node, + n_local, + ei, + ev, + em, + do, + drp, + so, + srp, + destination_sorted=True, + fparam=fp, + aparam=ap, + do_atomic_virial=True, + charge_spin=cs, + ) + # The graph AOTI artifact was frozen via forward_lower_graph_exportable, + # which translates the internal fitting keys ("energy_redu", + # "energy_derv_r", ...) to the public forward_lower convention + # ("energy", "force", ...) -- see ener_model.py:441. Apply the same + # translation to the eager reference so the two key sets line up. + eager_out = _translate_energy_keys( + eager_internal, + do_grad_r=model.do_grad_r("energy"), + do_grad_c=model.do_grad_c("energy"), + do_atomic_virial=True, + local=True, + ) + + # Anti-vacuity guard: fresh DPA4 is edge-independent (forces ~0); + # confirm the jitter above made the eager reference force + # non-trivial, else the AOTI parity below would pass trivially. + f_ref = eager_out["force"].detach().cpu().numpy() + assert np.abs(f_ref).max() > 1e-6, ( + f"eager reference force is near-zero ({np.abs(f_ref).max():.3e}); " + f"jitter not effective -- AOTI parity check would be vacuous" + ) + + # The artifact is compiled for _env.DEVICE (move_to_device_pass in + # deserialize_to_file); move inputs there while the eager reference + # above stays on CPU. + ( + d_atype, + d_n_node, + d_n_local, + d_ei, + d_ev, + d_em, + d_do, + d_drp, + d_so, + d_srp, + d_fp, + d_ap, + d_cs, + ) = _to_artifact_device( + atype, n_node, n_local, ei, ev, em, do, drp, so, srp, fp, ap, cs + ) + artifact_out = regular( + d_atype, + d_n_node, + d_n_local, + d_ei, + d_ev, + d_em, + d_do, + d_drp, + d_so, + d_srp, + d_fp, + d_ap, + d_cs, + ) + + # Compare every key the artifact produces against the (translated) + # eager reference. + compared = 0 + for key, val in artifact_out.items(): + if key not in eager_out or eager_out[key] is None or val is None: + continue + np.testing.assert_allclose( + val.detach().cpu().numpy(), + eager_out[key].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + err_msg=( + f"artifact vs eager forward_common_lower_graph differs: {key}" + ), + ) + compared += 1 + assert compared > 0, ( + f"no overlapping output keys compared; artifact keys=" + f"{sorted(artifact_out)}, eager keys={sorted(eager_out)}" + ) + assert "energy_redu" in artifact_out or "energy" in artifact_out + + +# ============================================================================= +# Task 6: graph-kind ``.pt2`` freeze for the NATIVE-spin DPA4 wrapper +# (``NativeSpinEnergyModel``, type ``native_spin``) -- spin rides the +# NeighborGraph lower ONLY (no dense/nlist lower, no with-comm sidecar: see +# ``_needs_with_comm_artifact``'s native-spin first rule). The VIRTUAL-atom +# spin scheme (``SpinModel``, type ``spin_ener``) has no graph-lower +# implementation at all and must keep raising ``NotImplementedError``. +# ============================================================================= + +# Minimal virtual-atom (deepspin) spin config: a non-GNN se_e2_a backbone is +# enough to prove the graph-form rejection still fires for "spin_ener" -- +# the rejection happens before any tracing/AOTI compile, so this stays fast. +_VIRTUAL_SPIN_CONFIG = { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": [4, 4], + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [4, 4], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "neuron": [4, 4], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + }, + "spin": { + "use_spin": [True, False], + "virtual_scale": [0.3140], + }, +} + + +def _freeze_native_spin(model_file) -> None: + """Freeze a jittered native-spin DPA4 model to a graph-kind ``.pt2``. + + Mirrors ``test_dpa4_freeze_to_pt2``'s serialize -> ``deserialize_to_file`` + sequence, but native spin has ONLY a graph lower (no dense/nlist lower -- + see ``NativeSpinEnergyModel``'s module docstring), so ``lower_kind="graph"`` + is explicit rather than parametrized. ``_build_native_spin_model_cpu`` + already jitters DPA4's zero-initialized residual projections (else the + model is architecturally spin-independent -- see that helper's + docstring), so no extra jitter step is needed here. + """ + model = _build_native_spin_model_cpu() + data = {"model": model.serialize()} + deserialize_to_file(str(model_file), data, lower_kind="graph") + + +@pytest.mark.skipif( + os.environ.get("CI") == "true", + reason="AOTInductor compile is slow (minutes); run locally only by default.", +) +def test_native_spin_graph_freeze(tmp_path) -> None: + """Native-spin DPA4 freezes to a graph-kind .pt2: metadata + no sidecar.""" + model_file = tmp_path / "dpa4_spin_graph.pt2" + _freeze_native_spin(model_file) + + with zipfile.ZipFile(model_file) as z: + names = z.namelist() + # AOTInductor also embeds its OWN internal per-kernel files whose + # names end with "metadata.json" (e.g. + # ".wrapper_metadata.json", ".kernel_metadata.json" + # under "data/aotinductor/model/") -- read the exact PyTorch + # PT2_EXTRA_PREFIX path for OUR sidecar, not a fragile + # ``endswith("metadata.json")`` scan (mirrors + # ``test_dpa4_freeze_to_pt2`` above). + md = json.loads(z.read("model/extra/metadata.json").decode("utf-8")) + + assert md["type_map"] == NATIVE_SPIN_CONFIG["type_map"] + assert md["lower_input_kind"] == "graph" + assert md["is_spin"] is True + # Native spin participates in the with-comm artifact on the GRAPH lower + # (pt's SeZMNativeSpinModel does not override supports_edge_parallel): + # the spin input is per-node with ghost rows delivered by the LAMMPS + # ``sp`` forward-comm, so only the per-block ghost FEATURE refresh needs + # border_op -- the same one the energy model drives. + assert md["has_comm_artifact"] is True + assert md["has_message_passing"] is True + assert md["ntypes_spin"] == 1 # use_spin=[True, False] + assert md["use_spin"] == [True, False] + assert "force_mag" in md["output_keys"] + for key in ("atom_energy", "energy", "force", "virial"): + assert key in md["output_keys"] + # ... and the nested artifact must actually BE there. This assertion was + # inverted while native spin was single-rank; leaving it that way made the + # test contradict its own has_comm_artifact check above, and it survived + # only because the CPU freeze fails earlier on an unrelated inductor bug. + assert any(n.endswith("forward_lower_with_comm.pt2") for n in names), ( + "has_comm_artifact is True but the nested with-comm .pt2 is missing" + ) + + +def test_native_spin_nlist_deserialize_rejected(tmp_path) -> None: + """``deserialize_to_file`` fails fast on ``native_spin`` + non-graph lower. + + Native spin has no dense/nlist lower; without this guard the default + ``lower_kind="nlist"`` used to fail deep inside tracing with an opaque + ``TypeError``. The guard raises before any tracing/compile, so this + stays fast. + """ + model = _build_native_spin_model_cpu() + data = {"model": model.serialize()} + with pytest.raises(ValueError, match="only the NeighborGraph lower"): + deserialize_to_file(str(tmp_path / "spin_nlist.pte"), data, lower_kind="nlist") + + +@pytest.mark.skipif( + os.environ.get("CI") == "true", + reason="AOTInductor compile is slow (minutes); run locally only by default.", +) +def test_native_spin_default_freeze_routes_to_graph(tmp_path) -> None: + """Public ``freeze()`` default path yields a graph-kind ``.pt2``. + + Freezes a native-spin checkpoint WITHOUT supplying ``lower_kind`` (and + with a suffixless output, as the CLI passes it): the public freeze layer + must resolve native spin to the graph lower BEFORE choosing the export + ABI and the default output suffix. + """ + import copy + + from deepmd.pt_expt.entrypoints.main import ( + freeze, + ) + from deepmd.pt_expt.train.wrapper import ( + ModelWrapper, + ) + + model = _build_native_spin_model_cpu() + wrapper = ModelWrapper(model, model_params=copy.deepcopy(NATIVE_SPIN_CONFIG)) + ckpt = tmp_path / "model.pt" + torch.save({"model": wrapper.state_dict()}, ckpt) + + output = tmp_path / "frozen_native_spin" # suffixless: default CLI form + freeze(model=str(ckpt), output=str(output)) + + pt2 = output.with_suffix(".pt2") + assert pt2.exists(), "default suffix must follow the resolved graph kind" + with zipfile.ZipFile(pt2) as z: + md = json.loads(z.read("model/extra/metadata.json").decode("utf-8")) + assert md["lower_input_kind"] == "graph" + assert md["is_spin"] is True + + +def test_lower_kind_override_is_logged(tmp_path, caplog) -> None: + """Overriding the caller's lower_kind must be WARNED, not silent. + + The dense lower is deprecated here, so a graph-capable model is forced + onto the graph lower even when the caller asked for "nlist". The two + lowers are not numerically identical, so the override has to be visible + in the log rather than inferred from the output suffix. + + The export itself is stubbed: this pins the resolution and its log, and + a real compile would only add minutes (and this workstation's inductor + bug) without testing anything extra. + """ + import copy + import logging + from unittest import ( + mock, + ) + + from deepmd.pt_expt.entrypoints.main import ( + freeze, + ) + from deepmd.pt_expt.model.get_model import ( + get_model, + ) + from deepmd.pt_expt.train.wrapper import ( + ModelWrapper, + ) + + config = copy.deepcopy(_DPA4_CONFIG) + model = get_model(copy.deepcopy(config)).to(torch.device("cpu")).eval() + wrapper = ModelWrapper(model, model_params=copy.deepcopy(config)) + ckpt = tmp_path / "model.pt" + torch.save({"model": wrapper.state_dict()}, ckpt) + + seen: dict = {} + + def _stub(model_file, data, *args, **kwargs): + seen["model_file"] = model_file + seen["lower_kind"] = kwargs.get("lower_kind") + + with ( + # ``freeze`` imports it inside the function, so patch it at source + mock.patch("deepmd.pt_expt.utils.serialization.deserialize_to_file", _stub), + caplog.at_level(logging.WARNING, logger="deepmd.pt_expt.entrypoints.main"), + ): + freeze(model=str(ckpt), output=str(tmp_path / "frozen"), lower_kind="nlist") + + warned = [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] + assert any("OVERRIDDEN" in m for m in warned), ( + f"the lower_kind override was not warned about; got {warned}" + ) + # ... and it really did override, both in the export call and the suffix + assert seen["lower_kind"] == "graph" + assert seen["model_file"].endswith(".pt2") + + +def test_bridged_default_freeze_routes_to_graph(tmp_path) -> None: + """A ZBL-bridged model also default-freezes to a graph-kind ``.pt2``. + + The dense lower is deprecated in this backend, and for a bridged model + it does not exist at all: the analytical term's ``forward_atomic`` + raises on the dense route, so the public default + ``lower_kind="nlist"`` used to fail deep inside the dense trace. The + resolution was previously an ``isinstance(NativeSpinModelKind)`` + special case, which a bridged model is not an instance of; it now asks + the graph-lower capability, which covers both. + """ + import copy + + from deepmd.pt_expt.entrypoints.main import ( + freeze, + ) + from deepmd.pt_expt.model.get_model import ( + get_model, + ) + from deepmd.pt_expt.train.wrapper import ( + ModelWrapper, + ) + + config = copy.deepcopy(_DPA4_CONFIG) + config["bridging_method"] = "ZBL" + config["bridging_r_inner"] = 0.8 + config["bridging_r_outer"] = 1.2 + model = get_model(copy.deepcopy(config)).to(torch.device("cpu")).eval() + wrapper = ModelWrapper(model, model_params=copy.deepcopy(config)) + ckpt = tmp_path / "model.pt" + torch.save({"model": wrapper.state_dict()}, ckpt) + + output = tmp_path / "frozen_bridged" # suffixless: default CLI form + freeze(model=str(ckpt), output=str(output)) + + pt2 = output.with_suffix(".pt2") + assert pt2.exists(), "default suffix must follow the resolved graph kind" + with zipfile.ZipFile(pt2) as z: + md = json.loads(z.read("model/extra/metadata.json").decode("utf-8")) + assert md["lower_input_kind"] == "graph" + + +def test_virtual_spin_graph_freeze_still_rejected(tmp_path) -> None: + """spin_ener (virtual) graph freeze keeps raising ``NotImplementedError``. + + The virtual-atom scheme doubles the atom count (real + virtual) and has + no graph-lower implementation; only the native scheme + (``native_spin``) is graph-eligible (see the module docstring + above). + """ + model = get_model(_VIRTUAL_SPIN_CONFIG) + model.to("cpu") + model.eval() + data = {"model": model.serialize()} + assert data["model"]["type"] == "spin_ener" + + with pytest.raises(NotImplementedError, match="graph-form"): + deserialize_to_file( + str(tmp_path / "virtual_spin_graph.pt2"), data, lower_kind="graph" ) - compared += 1 - # Guard against a vacuous pass (no overlapping keys compared). - assert compared > 0, ( - f"no overlapping output keys compared; artifact keys=" - f"{sorted(artifact_out)}, eager keys={sorted(eager_out)}" - ) - # The energy output must be among the compared keys. - assert "energy_redu" in artifact_out or "energy" in artifact_out + + +# ============================================================================= +# Task 7: DeepEval graph fast path for the NATIVE-spin ``.pt2`` -- the frozen +# artifact from Task 6 above must be evaluable through the public DeepPot API +# (``deepmd.pt_expt.infer.deep_eval.DeepEval._eval_model_graph_spin``), NOT +# just constructible. Compares against the EAGER pt_expt +# ``NativeSpinEnergyModel.forward`` on the SAME weights/system (rtol=atol=1e-10, +# CPU fp64 -- project convention for same-math weight-copied parity). +# ============================================================================= + +_SPIN_EVAL_NATOMS = 6 +_SPIN_EVAL_ATYPES = np.array([0, 0, 0, 1, 1, 1], dtype=np.int32) # Ni,Ni,Ni,O,O,O +_SPIN_EVAL_COORDS = np.array( + [ + [1.0, 1.0, 1.0], + [3.2, 1.4, 1.1], + [1.3, 1.8, 1.0], + [0.4, 1.2, 1.6], + [3.6, 2.0, 1.3], + [3.4, 0.7, 1.7], + ], + dtype=np.float64, +).reshape(1, _SPIN_EVAL_NATOMS, 3) +_SPIN_EVAL_CELL = (np.eye(3, dtype=np.float64) * 6.0).reshape(1, 9) +# Deliberately NOT pre-masked by type (mirrors TestNativeSpinEnergyModelPtExpt): +# the model's own descriptor gating must zero the non-spin (type 1) rows. +_SPIN_EVAL_SPINS = np.array( + [ + [0.11, 0.05, -0.02], + [-0.07, 0.09, 0.03], + [0.02, -0.06, 0.08], + [0.01, -0.01, 0.02], + [-0.02, 0.03, -0.01], + [0.015, 0.02, -0.03], + ], + dtype=np.float64, +).reshape(1, _SPIN_EVAL_NATOMS, 3) + + +@pytest.mark.skipif( + os.environ.get("CI") == "true", + reason="AOTInductor compile is slow (minutes); run locally only by default.", +) +def test_deep_eval_graph_spin_parity(tmp_path) -> None: + """DeepEval on a graph-kind native-spin ``.pt2`` matches eager ``forward``. + + Exercises ``DeepEval._eval_model_spin``'s ``lower_input_kind == "graph"`` + branch end to end through the public ``DeepPot`` API: energy, force, + force_mag and (global) virial must reproduce the eager + ``NativeSpinEnergyModel.forward`` on the identical weights and system. + ``mask_mag`` is not exported by the graph ABI; the DeepEval adapter + synthesizes it from the artifact's ``use_spin`` metadata and the input + atom types (see ``_eval_model_graph_spin``), and this test asserts the + exact boolean mask through the public DeepPot API. + """ + from deepmd.infer import ( + DeepPot, + ) + + model = _build_native_spin_model_cpu() + + coord_t = torch.tensor(_SPIN_EVAL_COORDS, dtype=torch.float64) + atype_t = torch.tensor( + _SPIN_EVAL_ATYPES.reshape(1, _SPIN_EVAL_NATOMS), dtype=torch.int64 + ) + spin_t = torch.tensor(_SPIN_EVAL_SPINS, dtype=torch.float64) + box_t = torch.tensor(_SPIN_EVAL_CELL, dtype=torch.float64) + ref = model.forward(coord_t, atype_t, spin_t, box=box_t) + + # Anti-vacuity: a bare (non-jittered) DPA4 zero-initializes residual + # projections, which would make force_mag identically zero and the + # parity check below vacuous by construction (see + # ``_build_native_spin_model_cpu``'s docstring for why jitter fixes + # this). + fm_max = ref["force_mag"].abs().max().item() + assert fm_max > 1e-6, ( + "expected the jittered model's force_mag to be non-trivial; got " + f"max |force_mag| = {fm_max:.3e} (jitter not effective -- the " + "parity check below would be vacuous)" + ) + + model_file = tmp_path / "dpa4_spin_graph_eval.pt2" + data = {"model": model.serialize()} + deserialize_to_file(str(model_file), data, lower_kind="graph") + + dp = DeepPot(str(model_file)) + assert dp.has_spin + e, f, v, fm, mm = dp.eval( + _SPIN_EVAL_COORDS, + _SPIN_EVAL_CELL, + _SPIN_EVAL_ATYPES, + atomic=False, + spin=_SPIN_EVAL_SPINS, + ) + + # Public DeepPot contract: mask_mag marks spin-active atoms + # (use_spin=[True, False] x atypes [0,0,0,1,1,1]). + np.testing.assert_array_equal( + mm.reshape(-1).astype(bool), + np.array([True, True, True, False, False, False]), + err_msg="mask_mag", + ) + np.testing.assert_allclose( + e.reshape(-1), + ref["energy"].detach().numpy().reshape(-1), + rtol=1e-10, + atol=1e-10, + err_msg="energy", + ) + np.testing.assert_allclose( + f.reshape(-1), + ref["force"].detach().numpy().reshape(-1), + rtol=1e-10, + atol=1e-10, + err_msg="force", + ) + np.testing.assert_allclose( + fm.reshape(-1), + ref["force_mag"].detach().numpy().reshape(-1), + rtol=1e-10, + atol=1e-10, + err_msg="force_mag", + ) + np.testing.assert_allclose( + v.reshape(-1), + ref["virial"].detach().numpy().reshape(-1), + rtol=1e-10, + atol=1e-10, + err_msg="virial", + ) + + +@pytest.mark.skipif( + os.environ.get("CI") == "true", + reason="AOTInductor compile is slow (minutes); run locally only by default.", +) +def test_native_spin_chg_spin_graph_freeze(tmp_path) -> None: + """COMBINED native-spin + charge-spin FiLM freezes to a graph .pt2. + + Review 3638047227: the combined public configuration must export. The + charge_spin slot rides the ABI tail (slot 13); metadata carries the + chg-spin fields the C++/DeepEval loaders key on. + """ + import copy + + from deepmd.pt_expt.descriptor.dpa4 import ( + DescrptDPA4, + ) + from deepmd.pt_expt.model.get_model import get_model as pt_expt_get_model + + from ...dpa4_fixtures import ( + jitter_zero_arrays, + ) + from .test_dpa4_native_spin import ( + COMBINED_CHG_SPIN_CONFIG, + ) + + cpu = torch.device("cpu") + model = pt_expt_get_model(copy.deepcopy(COMBINED_CHG_SPIN_CONFIG)) + ds = model.atomic_model.descriptor + jittered = jitter_zero_arrays(ds.serialize(), np.random.default_rng(21)) + model.atomic_model.descriptor = DescrptDPA4.deserialize(jittered).to(cpu) + model = model.to(cpu).eval() + data = {"model": model.serialize()} + model_file = tmp_path / "dpa4_spin_chg_graph.pt2" + deserialize_to_file(str(model_file), data, lower_kind="graph") + + with zipfile.ZipFile(model_file) as z: + md = json.loads(z.read("model/extra/metadata.json").decode("utf-8")) + assert md["is_spin"] is True + assert md["lower_input_kind"] == "graph" + assert md["has_chg_spin_ebd"] is True + assert md["dim_chg_spin"] == 2 + assert "force_mag" in md["output_keys"] + + # DeepEval through the compiled artifact: the slot-13 charge_spin is + # LIVE (an integer-valued change moves the energy) and matches the eager + # forward with the same conditioning at the cross-artifact tolerance. + from deepmd.infer import ( + DeepPot, + ) + + dp = DeepPot(str(model_file)) + assert dp.has_spin + cs1 = np.array([[1.0, 2.0]]) + e1, f1, _v1, fm1, _mm1 = dp.eval( + _SPIN_EVAL_COORDS, + _SPIN_EVAL_CELL, + _SPIN_EVAL_ATYPES, + atomic=False, + spin=_SPIN_EVAL_SPINS, + charge_spin=cs1, + ) + e0, _f0, _v0, _fm0, _mm0 = dp.eval( + _SPIN_EVAL_COORDS, + _SPIN_EVAL_CELL, + _SPIN_EVAL_ATYPES, + atomic=False, + spin=_SPIN_EVAL_SPINS, + charge_spin=np.array([[0.0, 0.0]]), + ) + de = float(np.abs(np.asarray(e1) - np.asarray(e0)).max()) + assert de > 1e-10, f"charge_spin slot dead through the artifact: {de:.3e}" + + coord_t = torch.tensor(_SPIN_EVAL_COORDS, dtype=torch.float64).reshape(1, -1, 3) + atype_t = torch.tensor(_SPIN_EVAL_ATYPES, dtype=torch.int64).reshape(1, -1) + box_t = torch.tensor(_SPIN_EVAL_CELL, dtype=torch.float64).reshape(1, 9) + spin_t = torch.tensor(_SPIN_EVAL_SPINS, dtype=torch.float64).reshape(1, -1, 3) + ref = model.forward( + coord_t, + atype_t, + spin_t, + box=box_t, + charge_spin=torch.tensor(cs1, dtype=torch.float64), + ) + np.testing.assert_allclose( + np.asarray(e1).reshape(-1), + ref["energy"].detach().numpy().reshape(-1), + rtol=1e-10, + atol=1e-10, + err_msg="combined energy vs eager", + ) + np.testing.assert_allclose( + np.asarray(fm1).reshape(-1), + ref["force_mag"].detach().numpy().reshape(-1), + rtol=1e-10, + atol=1e-10, + err_msg="combined force_mag vs eager", + ) diff --git a/source/tests/pt_expt/model/test_dpa4_graph_lower.py b/source/tests/pt_expt/model/test_dpa4_graph_lower.py new file mode 100644 index 0000000000..48788934bc --- /dev/null +++ b/source/tests/pt_expt/model/test_dpa4_graph_lower.py @@ -0,0 +1,713 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""DPA4 on the pt_expt NeighborGraph lower: routing, parity, persistence. + +The pt_expt model plumbing that routes ``forward_common`` through the +carry-all graph (default-flip ``_resolve_graph_method``, autograd force/ +virial via ``forward_common_lower_graph``, the persisted +``graph_lower_disabled`` escape hatch) is GENERIC: it keys off +``descriptor.uses_graph_lower()`` and was implemented once for dpa1/dpa2. +DPA4/SeZM inherits the routing with ZERO new plumbing code once +``DescrptDPA4.uses_graph_lower()`` reports ``True`` -- these tests PROVE +that inheritance (routing/parity/persistence), following the dpa2 pattern +in ``test_dpa2_graph_lower.py``. + +Task 8 (graph ``.pt2`` export) adds ``test_graph_lower_symbolic_trace`` and +``test_graph_lower_torch_export``: both went GREEN with zero production +changes -- ``forward_lower_graph_exportable``/``forward_common_lower_graph`` +and ``_build_graph_dynamic_shapes`` are already output-agnostic over the +descriptor, so DPA4 (channels=16, n_radial=8, lmax=2, mmax=1, n_blocks=2, +SO(3) grid readout) traces and ``torch.export``s through the SAME generic +machinery Task 7 built for dpa1/dpa2, with no ``int()``/``.item()`` calls or +data-dependent ``if`` on a traced tensor anywhere in the DPA4-specific code +path exercised here (``AOTI indirect-indexing`` was already globally +disabled for DPA4 per Task 6). No trap-class fix was needed. +""" + +import ctypes + +import numpy as np +import pytest +import torch + +from deepmd.pt_expt.descriptor.dpa4 import ( + DescrptDPA4, +) +from deepmd.pt_expt.model import ( + EnergyModel, +) +from deepmd.pt_expt.model.get_model import ( + get_model, +) +from deepmd.pt_expt.model.graph_lower import ( + model_uses_graph_lower, +) +from deepmd.pt_expt.utils import ( + env, +) + +from ...dpa4_fixtures import ( + jitter_zero_arrays, +) +from ...seed import ( + GLOBAL_SEED, +) + +# --------------------------------------------------------------------------- +# Self-comm-dict helper -- copied from +# ``source/tests/pt_expt/model/test_dpa2_graph_lower.py`` (also duplicated in +# ``source/tests/pt_expt/descriptor/test_repflow_parallel.py`` and +# ``test_repformer_parallel.py``; the repo precedent for this generic, +# model-independent helper is a per-file copy, not a cross-test-module +# import). Builds a single-rank, self-only MPI ``comm_dict`` whose effect is +# a plain gather from the sendlist-indexed local rows into the ghost slots, +# so the ``border_op`` self-send branch (no MPI runtime needed) can be +# exercised eagerly. + + +def _addr_of(np_arr: np.ndarray) -> int: + """Return the raw int address of a numpy array's data buffer.""" + return np_arr.ctypes.data_as(ctypes.c_void_p).value + + +def _build_self_comm_dict( + *, + nloc: int, + nghost: int, + sendlist_indices: np.ndarray, + keepalive: list, +) -> dict: + """Build a comm_dict for a single-rank self-exchange. + + ``sendlist_indices`` (int32, length ``nghost``) gives the local row to + copy into each successive ghost slot ``[nloc, nloc + nghost)``. Control + tensors are forced to CPU: the C++ ``border_op`` host-side code + dereferences ``data_ptr()`` directly. + """ + sendlist_indices = np.ascontiguousarray(sendlist_indices, dtype=np.int32) + keepalive.append(sendlist_indices) + addr = _addr_of(sendlist_indices) + return { + "send_list": torch.tensor([addr], dtype=torch.int64, device="cpu"), + "send_proc": torch.zeros(1, dtype=torch.int32, device="cpu"), + "recv_proc": torch.zeros(1, dtype=torch.int32, device="cpu"), + "send_num": torch.tensor([nghost], dtype=torch.int32, device="cpu"), + "recv_num": torch.tensor([nghost], dtype=torch.int32, device="cpu"), + "communicator": torch.zeros(1, dtype=torch.int64, device="cpu"), + "nlocal": torch.tensor(nloc, dtype=torch.int32, device="cpu"), + "nghost": torch.tensor(nghost, dtype=torch.int32, device="cpu"), + } + + +# Small fp64 DPA4/SeZM config -- copied verbatim from the descriptor block of +# ``source/tests/pt_expt/model/test_dpa4_export.py``'s ``_DPA4_CONFIG`` (that +# file documents it as "provably builds"; sel=20/rcut=4.0 is non-binding for +# the 6-atom fixture below -- the real max degree is 11, see the module +# docstring of ``test_dpa2_graph_lower.py`` for the non-binding-sel +# rationale). +_DPA4_CONFIG = { + "type": "dpa4", + "type_map": ["foo", "bar"], + "descriptor": { + "type": "dpa4", + "sel": 20, + "rcut": 4.0, + "channels": 16, + "n_radial": 8, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "precision": "float64", + "seed": 1, + }, + "fitting_net": { + "type": "dpa4_ener", + "neuron": [16], + "precision": "float64", + "seed": 1, + }, +} + + +def _make_model(device) -> EnergyModel: + """Build a graph-eligible pt_expt DPA4/SeZM model from the exported config.""" + model = get_model(_DPA4_CONFIG) + return model.to(device) + + +def _make_message_sensitive_model(device, seed: int = 99) -> EnergyModel: + """A ``_make_model()`` variant with the zero-init residuals jittered. + + DPA4 deliberately zero-initializes several residual output projections + (see the ``jitter_zero_arrays`` docstring in + ``source/tests/dpa4_fixtures.py``) so a freshly constructed, untrained + descriptor is architecturally edge/message INDEPENDENT: its scalar + read-out is exactly the type embedding regardless of geometry or + neighbors. That makes a bare ``_make_model()`` VACUOUS for a + graph-vs-dense parity check -- the two routes would agree trivially + because neither route's output depends on the edges at all. This + jitters every exactly-zero float array in the descriptor's serialized + parameter tree in place, so the model's energy genuinely depends on + the neighbor edges (pinned by an in-test coordinate-perturbation guard + in ``test_forward_common_graph_matches_dense``). + """ + model = _make_model(device) + ds = model.atomic_model.descriptor + data = ds.serialize() + data = jitter_zero_arrays(data, np.random.default_rng(seed)) + jittered = DescrptDPA4.deserialize(data).to(device) + # Standard torch.nn.Module submodule replacement: "descriptor" is + # already a registered submodule of atomic_model, so this rebinds the + # ``_modules`` entry in place. + model.atomic_model.descriptor = jittered + return model + + +def _small_graph_inputs( + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Small 6-atom fp64 fixture (same geometry recipe as + ``TestDpa4GraphLower.setup_method``) for descriptor-level graph tests + that don't need a full model forward. + """ + generator = torch.Generator(device=device).manual_seed(GLOBAL_SEED) + cell = torch.rand([3, 3], dtype=torch.float64, device=device, generator=generator) + cell = (cell + cell.T) + 5.0 * torch.eye(3, device=device) + coord = torch.rand([6, 3], dtype=torch.float64, device=device, generator=generator) + coord = torch.matmul(coord, cell).unsqueeze(0) # [1, 6, 3] + atype = torch.tensor([[0, 0, 0, 1, 1, 1]], dtype=torch.int64, device=device) + box = cell.reshape(1, 9) + return coord, atype, box + + +def _graph_from( + coord: torch.Tensor, atype: torch.Tensor, box: torch.Tensor, rcut: float +): + """Build a carry-all ``NeighborGraph`` for the small fixture.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, + ) + + return build_neighbor_graph(coord, atype, box, rcut) + + +def test_dpa4_exchange_border_op_self_communication() -> None: + """Real ``border_op`` through the self-send branch: owned rows must be + untouched; ghost rows must be overwritten with the sendlist-named owner + rows. Exercises the actual custom op end-to-end, no MPI runtime needed. + """ + from deepmd.pt_expt.descriptor.dpa4_nn.block import ( + exchange_ghost_features, + ) + from deepmd.pt_expt.utils.comm import ( + ensure_comm_registered, + ) + + ensure_comm_registered() + nlocal, nghost = 4, 3 + n, d, c = nlocal + nghost, 4, 5 + owners = np.array([1, 0, 2], dtype=np.int32) # ghost i mirrors owned row owners[i] + keepalive: list = [] + comm = _build_self_comm_dict( + nloc=nlocal, + nghost=nghost, + sendlist_indices=owners, + keepalive=keepalive, + ) + x = torch.arange(n * d * c, dtype=torch.float64).reshape(n, d, 1, c) + x_in = x.clone() + out = exchange_ghost_features(x, comm) + assert out.shape == (n, d, 1, c) + torch.testing.assert_close(out[:nlocal], x_in[:nlocal]) # owned untouched + for gi, owner in enumerate(owners): + torch.testing.assert_close(out[nlocal + gi], x_in[int(owner)]) + + +def test_dpa4_exchange_schedule_counts() -> None: + """Pin the per-block border-exchange schedule on OUR graph route + (mirrors the pt-native ``TestSeZMExchangeSchedule``, + ``source/tests/pt/model/test_sezm_parallel.py:467-529``): block 0 skips + the exchange unless ``use_env_seed``. + + COUNT DERIVATION: ``DescrptDPA4._block_comm`` + (``deepmd/dpmodel/descriptor/dpa4.py:2215``) forwards ``comm_dict`` to + every block except block 0 when ``use_env_seed`` is False -- the + fixture descriptor defaults ``use_env_seed=True`` (plan-premise + correction: block 0 forwards too here). Every DPA4 block + unconditionally runs exactly one SO(2) unit per forward (dpmodel + ``SeZMInteractionBlock.call`` invokes ``_run_so2_unit`` exactly once on + every branch -- fast path, full-attn-res, block-attn-res -- + ``deepmd/dpmodel/descriptor/dpa4_nn/block.py:840,894,958``), so the + expected count is derived from the descriptor's own ``blocks``/ + ``use_env_seed`` attributes below, not hardcoded (it would differ for a + fixture with an SO2-free block, which this one does not have). + """ + from deepmd.pt_expt.descriptor.dpa4_nn import block as blk_mod + + device = env.DEVICE + dd = _make_message_sensitive_model(device).atomic_model.descriptor + coord, atype, box = _small_graph_inputs(device) + graph = _graph_from(coord, atype, box, dd.get_rcut()) + + calls = {"n": 0} + + def _counting_exchange( + x: torch.Tensor, comm_dict: dict[str, torch.Tensor] + ) -> torch.Tensor: + calls["n"] += 1 + return x + + fake_comm = dict.fromkeys( + ( + "send_list", + "send_proc", + "recv_proc", + "send_num", + "recv_num", + "communicator", + "nlocal", + "nghost", + ) + ) + orig = blk_mod.exchange_ghost_features + blk_mod.exchange_ghost_features = _counting_exchange + try: + dd.call_graph(graph, atype.reshape(-1), comm_dict=fake_comm) + finally: + blk_mod.exchange_ghost_features = orig + + n_blocks = len(dd.blocks) + expected = n_blocks if dd.use_env_seed else n_blocks - 1 + assert calls["n"] == expected + + +def test_dpa4_exchange_rejects_spin_comm() -> None: + """A ``comm_dict`` describing a spin system raises ``NotImplementedError`` + -- spin models never route the graph lower (``disable_graph_lower``), + so a ``has_spin`` comm_dict reaching the graph exchange seam is a + programming error, not a supported configuration. + """ + from deepmd.pt_expt.descriptor.dpa4_nn.block import ( + exchange_ghost_features, + ) + + x = torch.zeros(3, 2, 1, 4, dtype=torch.float64) + with pytest.raises(NotImplementedError, match="spin"): + exchange_ghost_features(x, {"has_spin": torch.ones(1)}) + + +class TestDpa4GraphLower: + def setup_method(self) -> None: + self.device = env.DEVICE + self.natoms = 6 + self.nt = 2 + self.type_map = ["foo", "bar"] + + generator = torch.Generator(device=self.device).manual_seed(GLOBAL_SEED) + cell = torch.rand( + [3, 3], dtype=torch.float64, device=self.device, generator=generator + ) + cell = (cell + cell.T) + 5.0 * torch.eye(3, device=self.device) + self.cell = cell.unsqueeze(0) # [1, 3, 3] + coord = torch.rand( + [self.natoms, 3], + dtype=torch.float64, + device=self.device, + generator=generator, + ) + coord = torch.matmul(coord, cell) + self.coord = coord.unsqueeze(0).to(self.device) # [1, natoms, 3] + self.atype = torch.tensor( + [[0, 0, 0, 1, 1, 1]], dtype=torch.int64, device=self.device + ) + + def test_model_uses_graph_lower(self) -> None: + """A graph-eligible DPA4 model default-flips; the escape hatch flips + it back off. + """ + model = _make_model(self.device) + assert model_uses_graph_lower(model) is True + model.atomic_model.descriptor.disable_graph_lower() + assert model_uses_graph_lower(model) is False + + def test_graph_lower_disabled_buffer_roundtrip(self) -> None: + """``disable_graph_lower()`` is persisted state: it round-trips a + ``state_dict`` save/load (the ``graph_lower_disabled`` buffer), and a + PRE-knob checkpoint (no buffer key) still strict-loads, defaulting to + the graph route. + """ + model = _make_model(self.device) + model.atomic_model.descriptor.disable_graph_lower() + sd = model.state_dict() + assert any("graph_lower_disabled" in k for k in sd), ( + "the hatch must ride the state_dict" + ) + + fresh = _make_model(self.device) + assert fresh.atomic_model.descriptor.uses_graph_lower() is True + fresh.load_state_dict(sd) + assert fresh.atomic_model.descriptor.uses_graph_lower() is False, ( + "restored buffer must re-disable the graph route" + ) + + # back-compat: a checkpoint written before the knob was persisted + # lacks the buffer key -- the strict load must still succeed and + # keep the default (graph) route. + old_sd = { + k: v + for k, v in _make_model(self.device).state_dict().items() + if "graph_lower_disabled" not in k + } + fresh2 = _make_model(self.device) + fresh2.load_state_dict(old_sd) + assert fresh2.atomic_model.descriptor.uses_graph_lower() is True + + def test_forward_common_graph_matches_dense(self) -> None: + """Default-flip ``forward_common`` (graph) matches + ``neighbor_graph_method="legacy"`` (dense) at non-binding sel, on a + message-sensitive (jittered) model. + + Force tolerance is one decade looser than energy: the two routes + differentiate through different autograd chains (the graph's + edge-vec leaf vs the dense route's local-coordinate leaf). + """ + model = _make_message_sensitive_model(self.device) + model.eval() + box = self.cell.reshape(1, 9) + + graph = model.forward_common( + self.coord.clone().requires_grad_(True), self.atype, box + ) + dense = model.forward_common( + self.coord.clone().requires_grad_(True), + self.atype, + box, + neighbor_graph_method="legacy", + ) + tol_e = {"rtol": 1e-10, "atol": 1e-12} + tol_f = {"rtol": 1e-8, "atol": 1e-10} + torch.testing.assert_close(graph["energy_redu"], dense["energy_redu"], **tol_e) + torch.testing.assert_close( + graph["energy_derv_r"], dense["energy_derv_r"], **tol_f + ) + + # Anti-vacuity guard: a coordinate perturbation must move the + # energy. Without the jitter in _make_message_sensitive_model, DPA4's + # zero-initialized residual projections make the output exactly + # edge-independent, which would make the parity assertions above + # trivially (and meaninglessly) true. + perturbed_coord = self.coord.clone() + perturbed_coord[0, 0, 0] += 0.1 + perturbed = model.forward_common( + perturbed_coord.requires_grad_(True), self.atype, box + ) + e_diff = (perturbed["energy_redu"] - graph["energy_redu"]).abs().max().item() + assert e_diff > 1e-6, ( + f"expected the message-sensitive model's energy to depend on " + f"coordinates; got a change of only {e_diff:.3e} (jitter not " + f"effective -- parity check above would be vacuous)" + ) + + def test_graph_route_actually_taken(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Anti-vacuity: the default route genuinely calls + ``DescrptDPA4.call_graph``, and the disabled route does not (a + silent fallback to dense on both branches would make the routing + tests above vacuous). + """ + model = _make_model(self.device) + model.eval() + box = self.cell.reshape(1, 9) + + calls = {"n": 0} + original = DescrptDPA4.call_graph + + def _spy(self_, *args: object, **kwargs: object) -> object: + calls["n"] += 1 + return original(self_, *args, **kwargs) + + monkeypatch.setattr(DescrptDPA4, "call_graph", _spy) + + model.forward_common(self.coord.clone().requires_grad_(True), self.atype, box) + assert calls["n"] > 0, "default route must call DescrptDPA4.call_graph" + + calls["n"] = 0 + model.atomic_model.descriptor.disable_graph_lower() + model.forward_common(self.coord.clone().requires_grad_(True), self.atype, box) + assert calls["n"] == 0, "disabled route must not call DescrptDPA4.call_graph" + + def test_graph_lower_symbolic_trace(self) -> None: + """``make_fx`` symbolic trace of ``forward_lower_graph_exportable`` + reproduces the eager graph lower bit-tight, on a message-sensitive + (jittered) model. + + ``forward_common_lower_graph`` computes force/virial via a single + ``torch.autograd.grad`` backward through its own ``edge_vec`` leaf + (see ``edge_transform_output.py:106``); tracing the whole exportable + wrapper with ``make_fx`` therefore traces that backward pass too -- + this is what makes the DPA4 SO(3)/GridMLP compute graph, including + its analytic force/virial, ``.pt2``-exportable. ``model.to("cpu")`` + before tracing mirrors the real ``.pt2`` export path (make_fx traces + on CPU by design, ``serialization.py:924``) and the dpa1 CUDA lesson + (traced inputs and params must share a device). + """ + from deepmd.pt_expt.utils.serialization import ( + build_synthetic_graph_inputs, + ) + + model = _make_message_sensitive_model(self.device).to("cpu") + model.eval() + sample = build_synthetic_graph_inputs( + model, + e_max=175, + nframes=2, + nloc=7, + dtype=torch.float64, + device=torch.device("cpu"), + ) + atype, n_node, n_local, ei, ev, em, do, drp, so, srp, fp, ap, cs = sample + traced = model.forward_lower_graph_exportable( + atype, + n_node, + n_local, + ei, + ev, + em, + do, + drp, + so, + srp, + fparam=fp, + aparam=ap, + do_atomic_virial=True, + charge_spin=cs, + destination_sorted=True, + tracing_mode="symbolic", + _allow_non_fake_inputs=True, + ) + out = traced(atype, n_node, n_local, ei, ev, em, do, drp, so, srp, fp, ap, cs) + ref = model.forward_common_lower_graph( + atype, + n_node, + n_local, + ei, + ev, + em, + do, + drp, + so, + srp, + destination_sorted=True, + fparam=fp, + aparam=ap, + do_atomic_virial=True, + ) + for key in ("energy", "force", "virial"): + assert torch.isfinite(out[key]).all(), f"non-finite traced {key}" + tol = {"rtol": 1e-12, "atol": 1e-12} + torch.testing.assert_close(out["energy"], ref["energy_redu"], **tol) + torch.testing.assert_close( + out["force"], ref["energy_derv_r"].reshape(out["force"].shape), **tol + ) + torch.testing.assert_close( + out["virial"], ref["energy_derv_c_redu"].reshape(out["virial"].shape), **tol + ) + + def test_graph_lower_torch_export(self) -> None: + """``torch.export.export`` the traced graph lower with the + production dynamic shapes (``_build_graph_dynamic_shapes``): the + edge axis ``E``, the flat node axis ``N``, and the frame axis ``nf`` + are all dynamic. Must export without ``GuardOnDataDependentSymNode`` + and the resulting exported program must reproduce the eager graph + lower AND generalize to a different (smaller) system size than the + one it was traced/exported on -- proving the dynamism is real, not + an artifact baked to the trace-time shapes. + """ + from deepmd.pt_expt.utils.serialization import ( + _build_graph_dynamic_shapes, + build_synthetic_graph_inputs, + ) + + model = _make_message_sensitive_model(self.device).to("cpu") + model.eval() + sample = build_synthetic_graph_inputs( + model, + e_max=175, + nframes=2, + nloc=7, + dtype=torch.float64, + device=torch.device("cpu"), + ) + atype, n_node, n_local, ei, ev, em, do, drp, so, srp, fp, ap, cs = sample + traced = model.forward_lower_graph_exportable( + atype, + n_node, + n_local, + ei, + ev, + em, + do, + drp, + so, + srp, + fparam=fp, + aparam=ap, + do_atomic_virial=True, + charge_spin=cs, + destination_sorted=True, + tracing_mode="symbolic", + _allow_non_fake_inputs=True, + ) + dynamic_shapes = _build_graph_dynamic_shapes( + atype, n_node, n_local, ei, ev, em, do, drp, so, srp, fp, ap, cs + ) + exported = torch.export.export( + traced, + (atype, n_node, n_local, ei, ev, em, do, drp, so, srp, fp, ap, cs), + dynamic_shapes=dynamic_shapes, + strict=False, + prefer_deferred_runtime_asserts_over_guards=True, + ) + loaded = exported.module() + + # Re-run on a SMALLER system (different nframes/nloc/edge count) to + # prove the exported program is genuinely dynamic, not specialized + # to the trace-time shapes. + small = build_synthetic_graph_inputs( + model, + e_max=None, + nframes=1, + nloc=3, + dtype=torch.float64, + device=torch.device("cpu"), + ) + ( + s_atype, + s_n_node, + s_n_local, + s_ei, + s_ev, + s_em, + s_do, + s_drp, + s_so, + s_srp, + s_fp, + s_ap, + s_cs, + ) = small + out = loaded( + s_atype, + s_n_node, + s_n_local, + s_ei, + s_ev, + s_em, + s_do, + s_drp, + s_so, + s_srp, + s_fp, + s_ap, + s_cs, + ) + ref = model.forward_common_lower_graph( + s_atype, + s_n_node, + s_n_local, + s_ei, + s_ev, + s_em, + s_do, + s_drp, + s_so, + s_srp, + destination_sorted=True, + fparam=s_fp, + aparam=s_ap, + do_atomic_virial=True, + ) + for key in ("energy", "force", "virial"): + assert torch.isfinite(out[key]).all(), f"non-finite exported {key}" + tol = {"rtol": 1e-10, "atol": 1e-10} + torch.testing.assert_close(out["energy"], ref["energy_redu"], **tol) + torch.testing.assert_close( + out["force"], ref["energy_derv_r"].reshape(out["force"].shape), **tol + ) + torch.testing.assert_close( + out["virial"], ref["energy_derv_c_redu"].reshape(out["virial"].shape), **tol + ) + + def test_graph_lower_do_atomic_virial_filtering(self) -> None: + """``do_atomic_virial`` is a pure output filter, not a compute switch. + + Production semantics already compute the per-atom virial + unconditionally on the graph route and discard it when unrequested + (``edge_transform_output.py:143-146``; the freeze forces ``True`` + for graph kinds); ``_translate_energy_keys`` (``local=True``) then + emits ``atom_virial`` only when requested. This pins the ``False`` + filtering branch at the model-forward level, eagerly (no AOTI + compile, no ``make_fx`` trace): call ``forward_common_lower_graph`` + twice on the SAME jittered model and inputs -- once with + ``do_atomic_virial=True``, once ``False`` -- translate both to the + public keys, and assert the atom-virial key is the only asymmetry: + every key present in both outputs is bit-identical + (``torch.equal``), i.e. the flag must not perturb + energy/force/virial at all. + """ + from deepmd.pt_expt.model.ener_model import ( + _translate_energy_keys, + ) + from deepmd.pt_expt.utils.serialization import ( + build_synthetic_graph_inputs, + ) + + model = _make_message_sensitive_model(self.device).to("cpu") + model.eval() + sample = build_synthetic_graph_inputs( + model, + e_max=175, + nframes=2, + nloc=7, + dtype=torch.float64, + device=torch.device("cpu"), + ) + atype, n_node, n_local, ei, ev, em, do, drp, so, srp, fp, ap, cs = sample + + def _run(do_atomic_virial: bool) -> dict[str, torch.Tensor]: + model_ret = model.forward_common_lower_graph( + atype, + n_node, + n_local, + ei, + ev, + em, + do, + drp, + so, + srp, + destination_sorted=True, + fparam=fp, + aparam=ap, + do_atomic_virial=do_atomic_virial, + ) + return _translate_energy_keys( + model_ret, + do_grad_r=model.do_grad_r("energy"), + do_grad_c=model.do_grad_c("energy"), + do_atomic_virial=do_atomic_virial, + local=True, + ) + + out_true = _run(True) + out_false = _run(False) + + assert "atom_virial" in out_true, "requesting the flag must produce the key" + assert "atom_virial" not in out_false, ( + "not requesting the flag must drop the key" + ) + # Sanity: atom_virial is the ONLY asymmetry between the two outputs. + assert set(out_true) - {"atom_virial"} == set(out_false) + for key in out_false: + assert torch.equal(out_true[key], out_false[key]), ( + f"do_atomic_virial must be a pure output filter; {key} differs " + f"between the True/False calls" + ) diff --git a/source/tests/pt_expt/model/test_dpa4_interop.py b/source/tests/pt_expt/model/test_dpa4_interop.py index 15771f20a6..9f4829d1e1 100644 --- a/source/tests/pt_expt/model/test_dpa4_interop.py +++ b/source/tests/pt_expt/model/test_dpa4_interop.py @@ -20,6 +20,9 @@ import torch from deepmd.pt.model.model import get_model as pt_get_model +from deepmd.pt_expt.model.dpa4_model import ( + DPA4EnergyModel, +) from deepmd.pt_expt.model.ener_model import ( EnergyModel, ) @@ -110,11 +113,11 @@ def test_serialize_layout(self, pt_dpa4_model) -> None: """The pt serialize layout matches the interop override's expectations.""" ser = pt_dpa4_model.serialize() # wrapper: recognised model type + @version 1 - assert ser["type"].lower() in BaseModel._SEZM_MODEL_TYPES + assert ser["type"].lower() in ("sezm", "dpa4") assert ser["@version"] == 1 # nested atomic: sezm_atomic @version 3 carrying the pt-only dens state atomic = ser["atomic_model"] - assert atomic["type"] in BaseModel._SEZM_ATOMIC_TYPES + assert atomic["type"] == "sezm_atomic" assert atomic["@version"] == 3 assert "dens_force_rmsd" in atomic["@variables"] assert "active_mode" in atomic @@ -132,7 +135,7 @@ def test_variables_filtered_to_out_bias_out_std(self, pt_dpa4_model) -> None: """The pt-only ``dens_force_rmsd`` @variable is dropped on normalize.""" atomic = pt_dpa4_model.serialize()["atomic_model"] assert set(atomic["@variables"]) >= {"out_bias", "out_std", "dens_force_rmsd"} - normalized = BaseModel._normalize_pt_sezm_atomic(atomic) + normalized = DPA4EnergyModel._normalize_pt_sezm_atomic(atomic) assert set(normalized["@variables"]) == {"out_bias", "out_std"} # version coerced to the standard atomic schema, type rewritten assert normalized["@version"] == 2 @@ -203,7 +206,7 @@ def test_atomic_version_in_range_accepted(self, pt_dpa4_model, version) -> None: """Both in-range atomic @versions {2, 3} normalize without raising.""" atomic = pt_dpa4_model.serialize()["atomic_model"] atomic["@version"] = version - normalized = BaseModel._normalize_pt_sezm_atomic(atomic) + normalized = DPA4EnergyModel._normalize_pt_sezm_atomic(atomic) assert normalized["@version"] == 2 diff --git a/source/tests/pt_expt/model/test_dpa4_native_spin.py b/source/tests/pt_expt/model/test_dpa4_native_spin.py new file mode 100644 index 0000000000..e16a08172a --- /dev/null +++ b/source/tests/pt_expt/model/test_dpa4_native_spin.py @@ -0,0 +1,1332 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""``force_mag`` autograd on the DPA4 pt_expt NeighborGraph lower. + +Task 3 of the "DPA4 native spin on the NeighborGraph route" plan wires a +SECOND autograd leaf (``spin``) alongside the existing ``edge_vec`` leaf in +``forward_common_lower_graph``/``fit_output_to_model_output_graph``: every +``r_differentiable`` reducible output additionally emits +``_derv_r_mag = -d_redu/dspin``. This exercises the pt_expt +BACKBONE energy model directly (a plain "dpa4" model config with +``use_spin`` set on the descriptor) -- NOT the ``NativeSpinEnergyModel`` +wrapper (that is Task 4); ``get_sezm_model`` only rejects a top-level +``"spin"`` key, so setting ``use_spin`` on the descriptor of an otherwise +plain ``"dpa4"`` model config reaches this trunk directly via +``model.call_common(coord, atype, box, spin=...)``. +""" + +import copy +import os + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.train import ( + DEFAULT_TASK_KEY, +) +from deepmd.pt.model.model import get_model as pt_get_model +from deepmd.pt_expt.descriptor.dpa4 import ( + DescrptDPA4, +) +from deepmd.pt_expt.entrypoints.main import ( + get_trainer, +) +from deepmd.pt_expt.fitting.dpa4_ener import ( + SeZMEnergyFittingNet, +) +from deepmd.pt_expt.model import ( + EnergyModel, +) +from deepmd.pt_expt.model.get_model import ( + get_model, +) +from deepmd.pt_expt.model.native_spin_model import ( + NativeSpinEnergyModel, +) +from deepmd.pt_expt.utils import env as _env +from deepmd.utils.argcheck import ( + normalize, +) +from deepmd.utils.compat import ( + update_deepmd_input, +) + +from ...dpa4_fixtures import ( + jitter_zero_arrays, +) +from ...seed import ( + GLOBAL_SEED, +) + +# Small fp64 DPA4/SeZM config with native spin enabled on the descriptor +# (``use_spin=[True, False]``: type 0 ("foo") carries spin, type 1 ("bar") +# does not). No top-level "spin" key -> ``get_sezm_model`` builds the plain +# backbone ``EnergyModel``, not the ``NativeSpinEnergyModel`` wrapper. +_DPA4_SPIN_CONFIG = { + "type": "dpa4", + "type_map": ["foo", "bar"], + "descriptor": { + "type": "dpa4", + "sel": 20, + "rcut": 4.0, + "channels": 16, + "n_radial": 8, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "precision": "float64", + "seed": 1, + "use_spin": [True, False], + }, + "fitting_net": { + "type": "dpa4_ener", + "neuron": [16], + "precision": "float64", + "seed": 1, + }, +} + + +def _build_jittered_backbone(seed: int = 99) -> EnergyModel: + """Build the pt_expt DPA4 backbone with ``use_spin`` set, jittered. + + DPA4 deliberately zero-initializes several residual output projections + (see ``dpa4_fixtures.jitter_zero_arrays``), so a freshly constructed, + untrained descriptor is architecturally edge/message (and spin) + INDEPENDENT -- a bare model would make both the finite-difference and + neutrality checks below vacuous. Jittering the descriptor's zero-init + weight tree makes the energy genuinely depend on ``spin`` (verified + in-test by ``TestGraphForceMag.setup_method``'s anti-vacuity guard). + """ + model = get_model(_DPA4_SPIN_CONFIG) + ds = model.atomic_model.descriptor + data = ds.serialize() + data = jitter_zero_arrays(data, np.random.default_rng(seed)) + jittered = DescrptDPA4.deserialize(data).to(_env.DEVICE) + model.atomic_model.descriptor = jittered + return model.to(_env.DEVICE).eval() + + +def _finite_diff_mag(model_fn, spin: np.ndarray, eps: float = 1e-4) -> np.ndarray: + """Central finite difference of ``model_fn`` (scalar) w.r.t. every spin + component, with the ``force_mag = -dE/dspin`` sign convention baked in. + """ + fm = np.zeros_like(spin) + for i in np.ndindex(*spin.shape): + sp = spin.copy() + sp[i] += eps + ep = model_fn(sp) + sp = spin.copy() + sp[i] -= eps + em = model_fn(sp) + fm[i] = -(ep - em) / (2 * eps) + return fm + + +class TestGraphForceMag: + def setup_method(self) -> None: + self.device = _env.DEVICE + self.model = _build_jittered_backbone() + + generator = torch.Generator(device=self.device).manual_seed(GLOBAL_SEED) + natoms = 6 + cell = torch.rand( + [3, 3], dtype=torch.float64, device=self.device, generator=generator + ) + cell = (cell + cell.T) + 5.0 * torch.eye(3, device=self.device) + coord = torch.rand( + [natoms, 3], dtype=torch.float64, device=self.device, generator=generator + ) + coord = torch.matmul(coord, cell) + self.coord = coord.unsqueeze(0) # (1, natoms, 3) + self.atype = torch.tensor( + [[0, 0, 0, 1, 1, 1]], dtype=torch.int64, device=self.device + ) + self.box = cell.reshape(1, 9) + spin = torch.rand( + [1, natoms, 3], dtype=torch.float64, device=self.device, generator=generator + ) + # only type 0 ("foo", use_spin=True) carries a magnetic moment; the + # non-magnetic type's spin input is inert (mirrors NativeSpinEnergyModel's + # mask_mag convention, dpmodel test_dpa4_native_spin_model.py). + self.spin = spin * (self.atype == 0)[..., None].to(spin.dtype) + + # Anti-vacuity guard: with the jitter applied, the energy must + # actually depend on spin (else the FD test below would trivially + # pass with both sides at zero). + out0 = self.model.call_common(self.coord, self.atype, self.box, spin=self.spin) + out1 = self.model.call_common( + self.coord, self.atype, self.box, spin=2.0 * self.spin + ) + e_diff = (out1["energy_redu"] - out0["energy_redu"]).abs().max().item() + assert e_diff > 1e-6, ( + f"expected the jittered model's energy to depend on spin; got a " + f"change of only {e_diff:.3e} (jitter not effective -- the FD " + f"test below would be vacuous)" + ) + + def test_force_mag_matches_finite_difference(self) -> None: + """``energy_derv_r_mag`` from the graph autograd == -dE/dspin by + central finite difference (atol 1e-6). + """ + out = self.model.call_common( + self.coord, + self.atype, + self.box, + spin=self.spin, + do_atomic_virial=False, + ) + fm = out["energy_derv_r_mag"] + + def _energy(sp: np.ndarray) -> float: + sp_t = torch.as_tensor(sp, device=self.device, dtype=self.coord.dtype) + ret = self.model.call_common(self.coord, self.atype, self.box, spin=sp_t) + return float(ret["energy_redu"].sum().detach()) + + fd = _finite_diff_mag(_energy, self.spin.cpu().numpy()) + np.testing.assert_allclose( + fm.squeeze(-2).detach().cpu().numpy().reshape(fd.shape), + fd, + atol=1e-6, + ) + + def test_force_unchanged_by_spin_leaf_wiring(self) -> None: + """``call_common`` WITHOUT ``spin`` has no ``energy_derv_r_mag`` key, + and the spin-less forward is deterministic (the new ``spin is not + None`` branch is a true no-op when ``spin`` is not supplied). + """ + out0 = self.model.call_common(self.coord, self.atype, self.box) + assert "energy_derv_r_mag" not in out0 + out1 = self.model.call_common(self.coord, self.atype, self.box) + assert "energy_derv_r_mag" not in out1 + # The graph-route reduction (segment_sum) and force assembly + # (edge_force_virial) scatter through ``index_add``, whose atomicAdd is + # non-deterministic on CUDA (1-2 fp64 ULP run-to-run); on CPU it is + # exact. So "the spin-less branch is a no-op" is a bit-exact claim on + # CPU and an eval-determinism claim (~1e-10) on CUDA. See CLAUDE.md. + det_rtol, det_atol = (0.0, 0.0) if self.device.type == "cpu" else (1e-10, 1e-12) + torch.testing.assert_close( + out0["energy_redu"], out1["energy_redu"], rtol=det_rtol, atol=det_atol + ) + torch.testing.assert_close( + out0["energy_derv_r"], out1["energy_derv_r"], rtol=det_rtol, atol=det_atol + ) + + def test_dense_route_spin_raises(self) -> None: + """Model-level spin rides ONLY the NeighborGraph lower (mirrors + ``test_dpa4_native_spin_model.py::test_dense_route_spin_raises``). + """ + with pytest.raises(NotImplementedError, match="NeighborGraph"): + self.model.call_common( + self.coord, + self.atype, + self.box, + spin=self.spin, + neighbor_graph_method="legacy", + ) + + +# ============================================================================= +# Task 4: the ``NativeSpinEnergyModel`` wrapper (top-level "spin" key, scheme +# "native") -- public forward() keys/shapes and weight-copied parity vs the +# pt ``SeZMNativeSpinModel`` reference. +# ============================================================================= + +# Same shape as dpmodel's NATIVE_SPIN_CONFIG +# (source/tests/common/dpmodel/test_dpa4_native_spin_model.py), plus the +# top-level "type": "dpa4" that pt's ``get_model`` requires to dispatch into +# ``get_sezm_spin_model`` (dpmodel's dispatch keys off `data["spin"]["scheme"]` +# alone and does not need it). +NATIVE_SPIN_CONFIG = { + "type": "dpa4", + "type_map": ["Ni", "O"], + "descriptor": { + "type": "dpa4", + "rcut": 4.0, + "sel": 8, + "channels": 16, + "n_radial": 8, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "precision": "float64", + "seed": 7, + "random_gamma": False, + }, + "fitting_net": { + "type": "dpa4_ener", + "neuron": [8, 8], + "precision": "float64", + "seed": 7, + }, + "spin": {"use_spin": [True, False], "scheme": "native"}, +} + + +def _jittered_wrapper(seed: int = 11) -> NativeSpinEnergyModel: + """Build the pt_expt ``NativeSpinEnergyModel`` wrapper, jittered. + + Mirrors ``_build_jittered_backbone`` above and + ``test_dpa4_native_spin_model.py::_jittered_model``: DPA4 zero-initializes + residual projections, so jittering the descriptor's serialized weight + tree is needed to make the wrapper's ``force_mag`` genuinely non-trivial. + """ + model = get_model(NATIVE_SPIN_CONFIG) + ds = model.atomic_model.descriptor + data = ds.serialize() + data = jitter_zero_arrays(data, np.random.default_rng(seed)) + model.atomic_model.descriptor = DescrptDPA4.deserialize(data).to(_env.DEVICE) + return model.to(_env.DEVICE).eval() + + +class TestNativeSpinEnergyModelPtExpt: + """Public ``forward()`` contract of the pt_expt ``NativeSpinEnergyModel``.""" + + def setup_method(self) -> None: + self.device = _env.DEVICE + self.model = _jittered_wrapper(seed=11) + + generator = torch.Generator(device=self.device).manual_seed(GLOBAL_SEED) + self.nf, self.nloc = 1, 6 + cell = torch.rand( + [3, 3], dtype=torch.float64, device=self.device, generator=generator + ) + cell = (cell + cell.T) + 5.0 * torch.eye(3, device=self.device) + coord = torch.rand( + [self.nloc, 3], dtype=torch.float64, device=self.device, generator=generator + ) + coord = torch.matmul(coord, cell) + self.coord = coord.unsqueeze(0) + # use_spin=[True, False]: type 0 ("Ni") carries spin, type 1 ("O") does not. + self.atype = torch.tensor( + [[0, 0, 0, 1, 1, 1]], dtype=torch.int64, device=self.device + ) + self.box = cell.reshape(1, 9) + # Deliberately NOT pre-masked by type: the model's own mask_mag/gating + # must zero the non-spin rows internally (see test below). + self.spin = torch.rand( + [self.nf, self.nloc, 3], + dtype=torch.float64, + device=self.device, + generator=generator, + ) + + def test_forward_keys_and_mask(self) -> None: + out = self.model.forward(self.coord, self.atype, self.spin, box=self.box) + for k in ("energy", "atom_energy", "force", "force_mag", "virial", "mask_mag"): + assert k in out + assert out["force_mag"].shape == (self.nf, self.nloc, 3) + assert out["force"].shape == (self.nf, self.nloc, 3) + assert out["mask_mag"].shape == (self.nf, self.nloc, 1) + # mask_mag: True only for the spin-active type (0) + expect_mask = (self.atype == 0).unsqueeze(-1) + torch.testing.assert_close(out["mask_mag"], expect_mask, rtol=0, atol=0) + + def test_force_mag_zero_on_non_spin_types(self) -> None: + """Non-spin-type (``atype==1``) rows of ``force_mag`` are exactly + zero, even though ``self.spin`` feeds nonzero noise there -- the + descriptor gates the spin embedding by type (docstring in + ``native_spin_model.py``), so no re-masking is applied. + """ + out = self.model.forward(self.coord, self.atype, self.spin, box=self.box) + non_spin = self.atype == 1 + assert torch.all(out["force_mag"][non_spin] == 0) + # anti-vacuity: the spin-active rows must be genuinely nonzero + spin_active = self.atype == 0 + assert out["force_mag"][spin_active].abs().max().item() > 1e-6 + + def test_atomic_virial_key_present_when_requested(self) -> None: + out = self.model.forward( + self.coord, self.atype, self.spin, box=self.box, do_atomic_virial=True + ) + assert "atom_virial" in out + assert out["atom_virial"].shape == (self.nf, self.nloc, 9) + + +def _pt_native_spin_model(seed: int = 3): + """Build the pt (reference) ``SeZMNativeSpinModel``, jittered.""" + model = pt_get_model(copy.deepcopy(NATIVE_SPIN_CONFIG)).to(torch.float64) + ds = model.atomic_model.descriptor + data = ds.serialize() + data = jitter_zero_arrays(data, np.random.default_rng(seed)) + from deepmd.pt.model.descriptor.sezm import ( + DescrptSeZM, + ) + + model.atomic_model.descriptor = DescrptSeZM.deserialize(data).to(torch.float64) + return model.eval() + + +class TestNativeSpinEnergyModelParity: + """Weight-copied fp64 parity vs ``deepmd.pt``'s ``SeZMNativeSpinModel``. + + Mirrors the component-weight-copy mechanism used throughout + ``test_dpa4_dpmodel_parity.py`` (build the pt reference, ``.serialize()`` + its descriptor/fitting, ``.deserialize()`` the result into the pt_expt + counterpart), lifted to the full model level: pt's ``DescrptSeZM``/ + ``SeZMEnergyFittingNet`` serialize to the SAME backend-agnostic dpmodel + dict schema pt_expt's ``DescrptDPA4``/``SeZMEnergyFittingNet`` + deserialize -- both are ports of one canonical dpmodel architecture. + Both models run on CPU fp64 (project convention: same-math + weight-copied fp64 parity ~= rtol/atol 1e-12). + """ + + def setup_method(self) -> None: + from deepmd.pt.model.task.sezm_ener import ( + SeZMEnergyFittingNet as PtSeZMEnergyFittingNet, + ) + + cpu = torch.device("cpu") + + # --- pt reference (jittered) --- + self.pt_model = _pt_native_spin_model(seed=3).to(cpu) + + # --- pt_expt model: same architecture, weights copied from pt --- + pt_expt_model = get_model(NATIVE_SPIN_CONFIG) + atomic = pt_expt_model.atomic_model + atomic.descriptor = DescrptDPA4.deserialize( + self.pt_model.atomic_model.descriptor.serialize() + ) + atomic.fitting_net = SeZMEnergyFittingNet.deserialize( + self.pt_model.atomic_model.fitting_net.serialize() + ) + self.pt_expt_model = pt_expt_model.to(cpu).eval() + assert isinstance( + self.pt_model.atomic_model.fitting_net, PtSeZMEnergyFittingNet + ) + + generator = torch.Generator(device=cpu).manual_seed(GLOBAL_SEED + 1) + self.nf, self.nloc = 1, 6 + cell = torch.rand([3, 3], dtype=torch.float64, generator=generator) + cell = (cell + cell.T) + 5.0 * torch.eye(3) + coord = torch.rand([self.nloc, 3], dtype=torch.float64, generator=generator) + coord = torch.matmul(coord, cell) + self.coord = coord.unsqueeze(0) + self.atype = torch.tensor([[0, 0, 0, 1, 1, 1]], dtype=torch.int64) + self.box = cell.reshape(1, 9) + self.spin = torch.rand( + [self.nf, self.nloc, 3], dtype=torch.float64, generator=generator + ) + + def test_parity_vs_pt_native_spin_model(self) -> None: + out_pt = self.pt_model.forward(self.coord, self.atype, self.spin, box=self.box) + out_pte = self.pt_expt_model.forward( + self.coord, self.atype, self.spin, box=self.box + ) + + # Anti-vacuity: fresh DPA4 zero-initializes residual projections, so + # a bare model would make force_mag identically zero on both sides + # and the parity check below vacuous by construction. Guard that the + # jitter made the magnetic force genuinely non-trivial. + fm_max = out_pte["force_mag"].abs().max().item() + assert fm_max > 1e-6, ( + f"expected the jittered model's force_mag to be non-trivial; " + f"got max |force_mag| = {fm_max:.3e} (jitter not effective -- " + f"the parity check below would be vacuous)" + ) + + for key in ("energy", "force", "force_mag", "virial"): + torch.testing.assert_close( + out_pt[key], out_pte[key], rtol=1e-12, atol=1e-12, msg=key + ) + torch.testing.assert_close( + out_pt["mask_mag"], out_pte["mask_mag"], rtol=0, atol=0, msg="mask_mag" + ) + + +# ============================================================================= +# Task 5: ``forward_lower_graph_exportable`` -- the graph-spin ``.pt2`` +# exportable ABI (spin at positional index 10). ``make_fx``/``torch.export`` +# tracing is CPU-only by design (``serialization.py:924`` moves the model +# to CPU before tracing) -- both tests below build the model AND the sample +# inputs on CPU explicitly, mirroring ``test_dpa4_graph_lower.py``'s +# ``test_graph_lower_symbolic_trace``/``test_graph_lower_torch_export``. +# ============================================================================= + + +def _build_native_spin_model_cpu(seed: int = 21) -> NativeSpinEnergyModel: + """Build the jittered pt_expt ``NativeSpinEnergyModel`` wrapper on CPU. + + Mirrors ``_jittered_wrapper`` above but pinned to CPU from construction + (export tracing is CPU-only; the dpa1 CUDA lesson is that traced inputs + and params must share a device). DPA4 zero-initializes residual + projections, so jittering is needed to make ``force_mag`` genuinely + non-trivial -- otherwise the trace-vs-eager comparisons below would be + vacuous (both sides identically zero). + """ + cpu = torch.device("cpu") + model = get_model(NATIVE_SPIN_CONFIG) + ds = model.atomic_model.descriptor + data = ds.serialize() + data = jitter_zero_arrays(data, np.random.default_rng(seed)) + model.atomic_model.descriptor = DescrptDPA4.deserialize(data).to(cpu) + return model.to(cpu).eval() + + +def _build_spin_graph_sample( + model: NativeSpinEnergyModel, + *, + nframes: int = 2, + nloc: int = 7, + e_max: int | None = 175, +) -> tuple[torch.Tensor, ...]: + """Build a small CPU sample matching the graph-spin positional ABI. + + Reuses :func:`build_synthetic_graph_inputs` (canonicalized, + destination-major) for the shared ``NeighborGraph`` CSR block, then adds + a ``spin`` tensor at index 10 -- a small NON-ZERO sample (not + ``torch.zeros``: an all-zero spin leaf can hit degenerate branches in + the equivariant spin embedding) matching the flat node axis ``N`` that + ``atype`` shares. + + Returns + ------- + tuple + ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, + destination_order, destination_row_ptr, source_order, + source_row_ptr, spin, fparam, aparam)`` -- the exact positional + order of ``NativeSpinEnergyModel.forward_lower_graph_exportable``. + """ + from deepmd.pt_expt.utils.serialization import ( + build_synthetic_graph_inputs, + ) + + sample = build_synthetic_graph_inputs( + model, + e_max=e_max, + nframes=nframes, + nloc=nloc, + dtype=torch.float64, + device=torch.device("cpu"), + want_charge_spin=False, + ) + (atype, n_node, n_local, ei, ev, em, do, drp, so, srp, fp, ap, _cs) = sample + generator = torch.Generator(device="cpu").manual_seed(GLOBAL_SEED) + spin = 0.1 + torch.rand(atype.shape[0], 3, dtype=torch.float64, generator=generator) + return atype, n_node, n_local, ei, ev, em, do, drp, so, srp, spin, fp, ap, _cs + + +class TestDPA4NativeSpinGraphLowerExportable: + """``forward_lower_graph_exportable`` establishes the positional ``.pt2`` + ABI for graph-spin models (spin at index 10, conditional ``charge_spin`` + tail at slot 13, no with-comm variant) -- the C++/serialization seams + mirror this ABI. + """ + + def test_graph_lower_exportable_symbolic_trace(self) -> None: + """``make_fx`` traces the graph-spin closure on CPU; the traced + output matches the eager ``forward_common_lower_graph`` reference + (same graph inputs -- the same physical system) bit-tight, and + includes a genuinely non-trivial ``force_mag``. + """ + model = _build_native_spin_model_cpu() + sample = _build_spin_graph_sample(model) + atype, n_node, n_local, ei, ev, em, do, drp, so, srp, spin, fp, ap, cs = sample + + traced = model.forward_lower_graph_exportable( + atype, + n_node, + n_local, + ei, + ev, + em, + do, + drp, + so, + srp, + spin, + fparam=fp, + aparam=ap, + do_atomic_virial=True, + destination_sorted=True, + tracing_mode="symbolic", + _allow_non_fake_inputs=True, + ) + out = traced( + atype, n_node, n_local, ei, ev, em, do, drp, so, srp, spin, fp, ap, cs + ) + for key in ( + "atom_energy", + "energy", + "force", + "force_mag", + "virial", + "atom_virial", + ): + assert key in out, f"missing output key {key}" + assert torch.isfinite(out[key]).all(), f"non-finite traced {key}" + + # Anti-vacuity: the jittered model's force_mag must be non-trivial, + # else the parity check below would trivially pass with both sides + # at zero. + assert out["force_mag"].abs().max().item() > 1e-6 + + ref = model.forward_common_lower_graph( + atype, + n_node, + n_local, + ei, + ev, + em, + do, + drp, + so, + srp, + destination_sorted=True, + do_atomic_virial=True, + fparam=fp, + aparam=ap, + spin=spin, + ) + tol = {"rtol": 1e-12, "atol": 1e-12} + torch.testing.assert_close(out["atom_energy"], ref["energy"], **tol) + torch.testing.assert_close(out["energy"], ref["energy_redu"], **tol) + torch.testing.assert_close( + out["force"], ref["energy_derv_r"].reshape(out["force"].shape), **tol + ) + torch.testing.assert_close( + out["force_mag"], + ref["energy_derv_r_mag"].reshape(out["force_mag"].shape), + **tol, + ) + torch.testing.assert_close( + out["virial"], ref["energy_derv_c_redu"].reshape(out["virial"].shape), **tol + ) + torch.testing.assert_close( + out["atom_virial"], + ref["energy_derv_c"].reshape(out["atom_virial"].shape), + **tol, + ) + + def test_graph_lower_exportable_torch_export(self) -> None: + """``torch.export.export`` succeeds with a dynamic ``nedge`` (and + ``N``/``nframes``) axis; the exported program reproduces the eager + graph lower AND generalizes to a different (smaller) system size + than the one it was traced/exported on. + """ + model = _build_native_spin_model_cpu() + sample = _build_spin_graph_sample(model) + atype, n_node, n_local, ei, ev, em, do, drp, so, srp, spin, fp, ap, cs = sample + + traced = model.forward_lower_graph_exportable( + atype, + n_node, + n_local, + ei, + ev, + em, + do, + drp, + so, + srp, + spin, + fparam=fp, + aparam=ap, + do_atomic_virial=True, + destination_sorted=True, + tracing_mode="symbolic", + _allow_non_fake_inputs=True, + ) + + nframes_dim = torch.export.Dim("nframes", min=1) + n_node_total_dim = torch.export.Dim("n_node_total", min=1) + nedge_dim = torch.export.Dim("nedge", min=2) + dynamic_shapes = ( + {0: n_node_total_dim}, # atype + {0: nframes_dim}, # n_node + {0: nframes_dim}, # n_local + {1: nedge_dim}, # edge_index + {0: nedge_dim}, # edge_vec + {0: nedge_dim}, # edge_mask + {0: nedge_dim}, # destination_order + {0: n_node_total_dim + 1}, # destination_row_ptr + {0: nedge_dim}, # source_order + {0: n_node_total_dim + 1}, # source_row_ptr + {0: n_node_total_dim}, # spin -- shares atype's N axis + {0: nframes_dim} if fp is not None else None, # fparam + {0: n_node_total_dim} if ap is not None else None, # aparam + {0: nframes_dim} if cs is not None else None, # charge_spin + ) + exported = torch.export.export( + traced, + (atype, n_node, n_local, ei, ev, em, do, drp, so, srp, spin, fp, ap, cs), + dynamic_shapes=dynamic_shapes, + strict=False, + prefer_deferred_runtime_asserts_over_guards=True, + ) + loaded = exported.module() + + # Re-run on a SMALLER system to prove the exported program is + # genuinely dynamic, not specialized to the trace-time shapes. + small = _build_spin_graph_sample(model, nframes=1, nloc=3, e_max=None) + ( + s_atype, + s_n_node, + s_n_local, + s_ei, + s_ev, + s_em, + s_do, + s_drp, + s_so, + s_srp, + s_spin, + s_fp, + s_ap, + s_cs, + ) = small + out = loaded( + s_atype, + s_n_node, + s_n_local, + s_ei, + s_ev, + s_em, + s_do, + s_drp, + s_so, + s_srp, + s_spin, + s_fp, + s_ap, + s_cs, + ) + ref = model.forward_common_lower_graph( + s_atype, + s_n_node, + s_n_local, + s_ei, + s_ev, + s_em, + s_do, + s_drp, + s_so, + s_srp, + destination_sorted=True, + do_atomic_virial=True, + fparam=s_fp, + aparam=s_ap, + spin=s_spin, + ) + for key in ("energy", "force", "force_mag", "virial", "atom_virial"): + assert torch.isfinite(out[key]).all(), f"non-finite exported {key}" + tol = {"rtol": 1e-10, "atol": 1e-10} + torch.testing.assert_close(out["energy"], ref["energy_redu"], **tol) + torch.testing.assert_close( + out["force"], ref["energy_derv_r"].reshape(out["force"].shape), **tol + ) + torch.testing.assert_close( + out["force_mag"], + ref["energy_derv_r_mag"].reshape(out["force_mag"].shape), + **tol, + ) + torch.testing.assert_close( + out["virial"], ref["energy_derv_c_redu"].reshape(out["virial"].shape), **tol + ) + torch.testing.assert_close( + out["atom_virial"], + ref["energy_derv_c"].reshape(out["atom_virial"].shape), + **tol, + ) + + +# ============================================================================= +# Task 11: training smoke -- native-spin DPA4 through the real pt_expt +# trainer (data loading, ``ener_spin`` loss dispatch, ``ModelWrapper``, +# graph-route autograd ``force_mag``), not a hand-rolled forward/backward. +# ============================================================================= + +# Small fp64 native-spin DPA4/SeZM training config. Reuses the real NiO spin +# dataset (``source/tests/pt/NiO/data``, type_map ["Ni", "O"], 32 atoms: 16 +# Ni carrying a magnetic moment, 16 O not) that ``source/tests/pt/ +# test_finetune_spin.py``'s ``TestSpinFinetuneSeA`` already exercises for the +# (unrelated) virtual-atom scheme with ``rcut=4.0``, ``sel=[20, 20]`` -- rcut +# reused verbatim here since it is known to see a sane number of neighbors on +# this system; DPA4's ``sel`` is only an initial search-capacity hint (grows +# on demand, never truncates the energy path), so a single scalar sel=40 +# (>= the se_e2_a per-type total of 40) is a safe, generous starting point. +# The descriptor is intentionally tiny (channels=8, n_radial=4, lmax=1, +# mmax=1, n_blocks=1) to keep the smoke test fast: this test proves the +# training PLUMBING (data requirement, loss dispatch, wrapper spin +# threading, autograd gradient flow), not model quality. +_TRAIN_MODEL_CONFIG = { + "type": "dpa4", + "type_map": ["Ni", "O"], + "descriptor": { + "type": "dpa4", + "rcut": 4.0, + "sel": 40, + "channels": 8, + "n_radial": 4, + "lmax": 1, + "mmax": 1, + "n_blocks": 1, + "precision": "float64", + "seed": 1, + }, + "fitting_net": { + "type": "dpa4_ener", + "neuron": [8], + "precision": "float64", + "seed": 1, + }, + "spin": {"use_spin": [True, False], "scheme": "native"}, +} + + +def _make_train_config(data_dir: str, numb_steps: int = 2) -> dict: + """Build a minimal native-spin DPA4 training config pointing at *data_dir*. + + Loss prefactors mirror ``source/tests/pt/test_finetune_spin.py``'s + ``TestSpinFinetuneSeA.setUp`` (the reference pt spin-training config): + ``ener_spin`` with both real-force (``fr``) and magnetic-force (``fm``) + terms enabled, so a nonzero magnetic force error actually contributes to + the loss (and therefore to the backward pass that reaches the + spin-embedding parameters). + """ + return { + "model": copy.deepcopy(_TRAIN_MODEL_CONFIG), + "learning_rate": { + "type": "exp", + "decay_steps": 500, + "start_lr": 0.001, + "stop_lr": 3.51e-8, + }, + "loss": { + "type": "ener_spin", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_fr": 1000, + "limit_pref_fr": 1, + "start_pref_fm": 1000, + "limit_pref_fm": 1, + }, + "training": { + "training_data": {"systems": [data_dir], "batch_size": 1}, + "validation_data": { + "systems": [data_dir], + "batch_size": 1, + "numb_btch": 1, + }, + "numb_steps": numb_steps, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 1, + "save_freq": numb_steps, + }, + } + + +class TestDPA4NativeSpinTrainingSmoke: + """End-to-end training smoke for the native-spin DPA4 trainer. + + Exercises the FULL production training path -- data loading, the + ``ener_spin`` loss dispatch, ``ModelWrapper.forward``, and the + graph-route backbone's autograd ``force_mag`` -- on the real NiO spin + dataset, not a hand-rolled forward/backward. + + Writing this test surfaced three real gaps (all in ``deepmd/dpmodel`` or + ``deepmd/pt_expt``, none in the read-only ``deepmd/pt``), fixed in the + same commit as this test: + + 1. ``ModelWrapper.forward`` (``deepmd/pt_expt/train/wrapper.py``) had no + ``spin`` parameter, so a spin-labeled batch's ``spin`` key made + ``self.wrapper(**input_dict, ...)`` raise ``TypeError``. Fixed by + mirroring ``deepmd.pt.train.wrapper.ModelWrapper.forward``'s + ``has_spin``-gated threading. + 2. ``get_additional_data_requirement`` + (``deepmd/pt_expt/train/training.py``) never declared ``spin`` as a + data requirement, so the data loader never learned to read + ``spin.npy`` in the first place. Fixed by mirroring + ``deepmd.pt.train.training.get_additional_data_requirement``'s + ``has_spin`` branch. + 3. ``NativeSpinEnergyModel.forward`` (``deepmd/pt_expt/model/ + dpa4_native_spin_model.py``) accepted no ``charge_spin`` keyword, but + ``ModelWrapper`` always forwards one -- fixed by accepting (and + passing through) ``charge_spin``, mirroring pt's + ``SeZMNativeSpinModel.forward``. Separately, ``EnergySpinLoss.call`` + (``deepmd/dpmodel/loss/ener_spin.py``) diffed the flat + ``(nf, natoms * 3)`` data-loader label directly against the model's + ``(nf, natoms, 3)`` prediction for ``force``/``force_mag`` -- unlike + every sibling atomic-label loss (``dos.py``/``tensor.py``), which + reshape the label to the canonical atomic shape before use. Fixed by + adding the same reshape. + """ + + def setup_method(self) -> None: + self.data_dir = os.path.join( + os.path.dirname(__file__), "..", "..", "pt", "NiO", "data", "single" + ) + if not os.path.isdir(self.data_dir): + pytest.skip(f"NiO spin data not found: {self.data_dir}") + + def test_training_smoke(self, tmp_path) -> None: + """Run 2 training steps; assert finite loss and a live spin-embedding grad.""" + config = _make_train_config(self.data_dir, numb_steps=2) + config = update_deepmd_input(config, warning=False) + config = normalize(config) + + old_cwd = os.getcwd() + os.chdir(tmp_path) + try: + trainer = get_trainer(config) + assert isinstance( + trainer.wrapper.model[DEFAULT_TASK_KEY], NativeSpinEnergyModel + ) + + tasks = trainer._make_training_tasks() + task = trainer.select_task(tasks) + + for step in range(2): + result = trainer.train_step(task, step) + loss = result.payload["loss"] + assert torch.isfinite(loss).all(), f"non-finite loss at step {step}" + + # force_mag gradients flow into training: after the last step's + # backward(), a spin-embedding parameter carries a nonzero grad. + # ``optimizer.step()`` reads (but does not clear) ``.grad``, so + # this checks the SAME gradient the optimizer just consumed -- + # ``train_step`` only calls ``zero_grad()`` at the START of the + # NEXT call, not after ``optimizer.step()``. + spin_params = [ + (name, p) + for name, p in trainer.wrapper.named_parameters() + if "spin_embedding" in name + ] + assert spin_params, "no spin_embedding parameter found in the model" + nonzero = [ + name + for name, p in spin_params + if p.grad is not None and torch.any(p.grad != 0) + ] + assert nonzero, ( + "no spin_embedding parameter has a nonzero grad after " + "training -- force_mag gradients are not flowing into " + "training" + ) + finally: + os.chdir(old_cwd) + + +class TestNativeSpinConfigFormsPtExpt: + """pt_expt twin of ``test_dpa4_native_spin_model.py::TestNativeSpinConfigForms``. + + ``spin.use_spin`` index/symbol forms are expanded against ``type_map`` + (``normalize_spin_use_spin``), and ``allow_missing_label`` is forwarded + into the constructed :class:`Spin` -- observable through the trainer's + spin data requirement. + """ + + @pytest.mark.parametrize( + ("use_spin_form", "expected"), + [ + (["Ni"], [True, False]), # element-symbol form + ([0], [True, False]), # type-index form + ([0, 1], [True, True]), # multiple type indices + ([True, False], [True, False]), # canonical boolean passthrough + ], + ) + def test_use_spin_forms(self, use_spin_form, expected) -> None: + config = copy.deepcopy(NATIVE_SPIN_CONFIG) + config["spin"] = {"use_spin": use_spin_form, "scheme": "native"} + model = get_model(config) + assert model.spin.use_spin.tolist() == expected + descriptor = model.atomic_model.descriptor + assert [bool(flag) for flag in descriptor.use_spin] == expected + + def test_use_spin_unknown_symbol_raises(self) -> None: + config = copy.deepcopy(NATIVE_SPIN_CONFIG) + config["spin"] = {"use_spin": ["Fe"], "scheme": "native"} + with pytest.raises(ValueError, match="absent from type_map"): + get_model(config) + + @pytest.mark.parametrize( + ("allow_missing", "expected_must"), + [ + (True, False), # relaxed: spin file optional, zero default + (False, True), # default: spin file mandatory + ], + ) + def test_allow_missing_label_data_requirement( + self, allow_missing, expected_must + ) -> None: + from deepmd.pt_expt.train.training import ( + get_additional_data_requirement, + ) + + config = copy.deepcopy(NATIVE_SPIN_CONFIG) + if allow_missing: + config["spin"]["allow_missing_label"] = True + model = get_model(config) + assert model.spin.allow_missing_label is allow_missing + reqs = get_additional_data_requirement(model) + spin_req = next(rr for rr in reqs if rr.key == "spin") + assert spin_req.must is expected_must + assert spin_req.default == 0.0 + + +class TestPublicBaseModelRoundTrip: + """Review 3638137290: pt_expt must round-trip its own native-spin + serialization through the PUBLIC BaseModel entry point, returning the + pt_expt class (a torch.nn.Module with .to() and the graph-export + machinery), with forward parity. + + Before the registry dispatch, ``BaseModel.deserialize`` entered a + hard-coded dpmodel branch and returned the numpy class here -- no + ``.to()`` (the concrete finetune failure: ``training.py`` calls + ``BaseModel.deserialize(...).to(DEVICE)``), no torch export machinery. + """ + + def test_roundtrip_returns_pt_expt_module(self) -> None: + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + + model = _jittered_wrapper(seed=11) + m2 = BaseModel.deserialize(model.serialize()) + assert type(m2) is NativeSpinEnergyModel + assert isinstance(m2, torch.nn.Module) + m2 = m2.to(_env.DEVICE) # the concrete finetune failure: .to(DEVICE) + assert hasattr(m2, "forward_common_lower_graph_exportable") + m2 = m2.eval() + generator = torch.Generator(device=_env.DEVICE).manual_seed(GLOBAL_SEED) + cell = torch.rand( + [3, 3], dtype=torch.float64, device=_env.DEVICE, generator=generator + ) + cell = (cell + cell.T) + 5.0 * torch.eye(3, device=_env.DEVICE) + coord = torch.matmul( + torch.rand( + [6, 3], dtype=torch.float64, device=_env.DEVICE, generator=generator + ), + cell, + ).unsqueeze(0) + atype = torch.tensor([[0, 0, 1, 0, 1, 1]], device=_env.DEVICE) + spin = torch.rand( + [1, 6, 3], dtype=torch.float64, device=_env.DEVICE, generator=generator + ) + box = cell.unsqueeze(0) + r1 = model(coord, atype, spin, box=box) + r2 = m2(coord, atype, spin, box=box) + for key in ("energy", "force", "force_mag"): + torch.testing.assert_close(r1[key], r2[key], rtol=1e-12, atol=1e-12) + + +# ============================================================================= +# Combined native spin + charge-spin FiLM (review 3638047227): pt does not +# reject this public configuration -- SeZMNativeSpinModel.forward accepts +# ``charge_spin`` alongside the native ``spin`` input. +# ============================================================================= + +COMBINED_CHG_SPIN_CONFIG = copy.deepcopy(NATIVE_SPIN_CONFIG) +COMBINED_CHG_SPIN_CONFIG["descriptor"]["add_chg_spin_ebd"] = True + + +class TestCombinedChargeSpinNativeSpin: + """Combined config builds, conditions on BOTH inputs, and matches pt. + + Weight-copied fp64 parity mechanics identical to + :class:`TestNativeSpinEnergyModelParity`. ``charge_spin`` is CATEGORICAL + (``ChargeSpinEmbedding`` casts the frame ``(charge, spin)`` pair to + int64 lookup indices), so the probe uses integer-valued changes. + """ + + def setup_method(self) -> None: + cpu = torch.device("cpu") + pt_model = pt_get_model(copy.deepcopy(COMBINED_CHG_SPIN_CONFIG)).to( + torch.float64 + ) + ds = pt_model.atomic_model.descriptor + data = jitter_zero_arrays(ds.serialize(), np.random.default_rng(3)) + from deepmd.pt.model.descriptor.sezm import ( + DescrptSeZM, + ) + + pt_model.atomic_model.descriptor = DescrptSeZM.deserialize(data).to( + torch.float64 + ) + self.pt_model = pt_model.eval().to(cpu) + + pt_expt_model = get_model(COMBINED_CHG_SPIN_CONFIG) + atomic = pt_expt_model.atomic_model + atomic.descriptor = DescrptDPA4.deserialize( + self.pt_model.atomic_model.descriptor.serialize() + ) + from deepmd.pt_expt.fitting.dpa4_ener import SeZMEnergyFittingNet as _FT + + atomic.fitting_net = _FT.deserialize( + self.pt_model.atomic_model.fitting_net.serialize() + ) + self.pt_expt_model = pt_expt_model.to(cpu).eval() + assert self.pt_expt_model.has_chg_spin_ebd() + assert self.pt_expt_model.has_spin() + + generator = torch.Generator(device=cpu).manual_seed(GLOBAL_SEED + 1) + self.nf, self.nloc = 1, 6 + cell = torch.rand([3, 3], dtype=torch.float64, generator=generator) + cell = (cell + cell.T) + 5.0 * torch.eye(3) + coord = torch.rand([self.nloc, 3], dtype=torch.float64, generator=generator) + self.coord = torch.matmul(coord, cell).unsqueeze(0) + self.atype = torch.tensor([[0, 0, 0, 1, 1, 1]], dtype=torch.int64) + self.box = cell.reshape(1, 9) + self.spin = torch.rand( + [self.nf, self.nloc, 3], dtype=torch.float64, generator=generator + ) + self.cs0 = torch.zeros([self.nf, 2], dtype=torch.float64) + self.cs1 = torch.tensor([[1.0, 2.0]], dtype=torch.float64) + + def test_combined_parity_and_sensitivity(self) -> None: + outs = {} + for tag, cs in (("cs0", self.cs0), ("cs1", self.cs1)): + out_pt = self.pt_model.forward( + self.coord, self.atype, self.spin, box=self.box, charge_spin=cs + ) + out_pte = self.pt_expt_model.forward( + self.coord, self.atype, self.spin, box=self.box, charge_spin=cs + ) + for key in ("energy", "force", "force_mag"): + torch.testing.assert_close( + out_pt[key], + out_pte[key], + rtol=1e-12, + atol=1e-12, + msg=f"{key}@{tag}", + ) + outs[tag] = out_pte + # charge_spin conditions the energy (reviewer's probe: nonzero + # response), and BOTH backends agree on the conditioned values above. + de = (outs["cs1"]["energy"] - outs["cs0"]["energy"]).abs().max().item() + assert de > 1e-10, f"charge_spin response vanished: {de:.3e}" + # native spin still conditions the same combined model + out_spin = self.pt_expt_model.forward( + self.coord, self.atype, 2.0 * self.spin, box=self.box, charge_spin=self.cs0 + ) + ds_ = (out_spin["energy"] - outs["cs0"]["energy"]).abs().max().item() + assert ds_ > 1e-10, f"spin response vanished in combined model: {ds_:.3e}" + + +class TestCombinedChargeSpinTrainingSmoke: + """Trainer smoke for the COMBINED native-spin + charge-spin FiLM model. + + NiO spin data carries no ``charge_spin`` file, so the run exercises the + ``default_chg_spin`` metadata path: ``get_additional_data_requirement`` + marks the ``charge_spin`` requirement optional with the descriptor's + default, and the FiLM conditions every batch on it. + """ + + def setup_method(self) -> None: + self.data_dir = os.path.join( + os.path.dirname(__file__), "..", "..", "pt", "NiO", "data", "single" + ) + if not os.path.isdir(self.data_dir): + pytest.skip(f"NiO spin data not found: {self.data_dir}") + + def test_training_smoke_combined(self, tmp_path) -> None: + config = _make_train_config(self.data_dir, numb_steps=2) + config["model"]["descriptor"]["add_chg_spin_ebd"] = True + config["model"]["descriptor"]["default_chg_spin"] = [0.0, 2.0] + config = update_deepmd_input(config, warning=False) + config = normalize(config) + + old_cwd = os.getcwd() + os.chdir(tmp_path) + try: + trainer = get_trainer(config) + model = trainer.wrapper.model[DEFAULT_TASK_KEY] + assert isinstance(model, NativeSpinEnergyModel) + assert model.has_chg_spin_ebd() + assert model.has_default_chg_spin() + + tasks = trainer._make_training_tasks() + task = trainer.select_task(tasks) + for step in range(2): + result = trainer.train_step(task, step) + loss = result.payload["loss"] + assert torch.isfinite(loss).all(), f"non-finite loss at step {step}" + finally: + os.chdir(old_cwd) + + +class TestCombinedChargeSpinGraphExportable: + """Combined native-spin + charge-spin FiLM through the graph exportable. + + ``charge_spin`` occupies the conditional slot-13 tail of the graph-spin + ``.pt2`` ABI; the traced module must (a) match the eager + ``forward_common_lower_graph`` with the SAME charge_spin, and (b) + respond to an integer-valued charge_spin change (the FiLM lookup is + categorical). + """ + + def test_combined_symbolic_trace_parity_and_sensitivity(self) -> None: + cpu = torch.device("cpu") + model = get_model(COMBINED_CHG_SPIN_CONFIG) + ds = model.atomic_model.descriptor + data = jitter_zero_arrays(ds.serialize(), np.random.default_rng(21)) + model.atomic_model.descriptor = DescrptDPA4.deserialize(data).to(cpu) + model = model.to(cpu).eval() + assert model.has_chg_spin_ebd() + + sample = _build_spin_graph_sample(model) + atype, n_node, n_local, ei, ev, em, do, drp, so, srp, spin, fp, ap, _ = sample + cs = torch.tensor([[1.0, 2.0]] * int(n_node.shape[0]), dtype=torch.float64) + + traced = model.forward_lower_graph_exportable( + atype, + n_node, + n_local, + ei, + ev, + em, + do, + drp, + so, + srp, + spin, + fparam=fp, + aparam=ap, + charge_spin=cs, + destination_sorted=True, + tracing_mode="symbolic", + _allow_non_fake_inputs=True, + ) + out = traced( + atype, n_node, n_local, ei, ev, em, do, drp, so, srp, spin, fp, ap, cs + ) + ref = model.forward_common_lower_graph( + atype, + n_node, + n_local, + ei, + ev, + em, + do, + drp, + so, + srp, + spin=spin, + charge_spin=cs, + destination_sorted=True, + ) + torch.testing.assert_close( + out["energy"], ref["energy_redu"], rtol=1e-12, atol=1e-12 + ) + torch.testing.assert_close( + out["force_mag"], + ref["energy_derv_r_mag"].squeeze(-2), + rtol=1e-12, + atol=1e-12, + ) + # charge_spin conditions the TRACED module (slot 13 is live, not a + # baked constant): an integer-valued change moves the energy. + cs0 = torch.zeros_like(cs) + out0 = traced( + atype, n_node, n_local, ei, ev, em, do, drp, so, srp, spin, fp, ap, cs0 + ) + de = (out["energy"] - out0["energy"]).abs().max().item() + assert de > 1e-10, f"charge_spin slot appears baked/dead: {de:.3e}" + + +class TestNativeSpinModelPairExcludeContract: + """Model-level ``pair_exclude_types`` on the native-spin graph route. + + Exclusion is a BUILD-time transform owned by the neighbor-graph + construction (decision #18/A4): the exported lower consumes a pre-excluded + ``edge_mask``, so every external feeder (Python ``DeepEval``, C++ + ``DeepSpinPTExpt::compute``) must fold it in from the + ``pair_exclude_types`` metadata field. These tests pin the two properties + the C++ regression + (``source/api_cc/tests/test_deepspin_dpa4_pairexcl_ptexpt.cc``) depends on: + the field reaches the metadata, and nothing inside the artifact reproduces + the mask. + """ + + @staticmethod + def _generic_model(pair_exclude_types: list | None, jitter_seed: int = 11): + """Build the pt_expt native-spin model the way the .pt2 fixture does. + + dpmodel builder -> serialize (+ inject the model-level exclusion) -> + pt_expt deserialize, mirroring ``source/tests/infer/gen_dpa4_spin.py``. + The generic (no ``type: dpa4``) config is what keeps + ``descriptor.exclude_types`` empty; the DPA4 alias would mirror the + pairs into it (see ``test_dpa4_alias_mirrors_into_descriptor``). + """ + from deepmd.dpmodel.model.model import get_model as dp_get_model + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + + config = copy.deepcopy(NATIVE_SPIN_CONFIG) + config.pop("type") + data = dp_get_model(config).serialize() + # DPA4 zero-initializes its residual projections, so an unjittered + # model is edge-independent and every exclusion is invisible. + data = jitter_zero_arrays(data, np.random.default_rng(jitter_seed)) + if pair_exclude_types is not None: + data["pair_exclude_types"] = copy.deepcopy(pair_exclude_types) + return BaseModel.deserialize(data).to(_env.DEVICE).eval() + + def test_generic_builder_keeps_descriptor_exclusions_empty(self) -> None: + """Negative contract: the mask must NOT be baked into the descriptor. + + A descriptor-level copy is compiled INTO the artifact and would make a + dead external seam invisible -- exactly the failure mode this pins. + """ + model = self._generic_model([[0, 1]]) + assert [list(p) for p in model.atomic_model.pair_exclude_types] == [[0, 1]] + assert not (model.atomic_model.descriptor.exclude_types or []) + + def test_dpa4_alias_mirrors_into_descriptor(self) -> None: + """Contrast: the ``type="dpa4"`` alias DOES mirror the pairs. + + Documented here so nobody rebuilds the C++ fixture from this config: + with the mirror in place the exclusion is applied inside the compiled + artifact and the external-seam regression becomes vacuous. + """ + config = copy.deepcopy(NATIVE_SPIN_CONFIG) # keeps type="dpa4" + config["pair_exclude_types"] = [[0, 1]] + model = get_model(config) + assert [list(p) for p in model.atomic_model.descriptor.exclude_types] == [ + [0, 1] + ] + + def test_metadata_carries_pair_exclude_types(self) -> None: + """The graph-lower metadata is what external feeders rebuild from.""" + from deepmd.pt_expt.utils.serialization import ( + _collect_metadata, + ) + + meta = _collect_metadata( + self._generic_model([[0, 1]]), is_spin=True, lower_kind="graph" + ) + assert meta["pair_exclude_types"] == [[0, 1]] + + base_meta = _collect_metadata( + self._generic_model(None), is_spin=True, lower_kind="graph" + ) + assert base_meta["pair_exclude_types"] == [] + + def test_exclusion_changes_the_prediction(self) -> None: + """Anti-vacuity: with identical weights, the mask must move the energy.""" + excluded = self._generic_model([[0, 1]]) + baseline = self._generic_model(None) + + generator = torch.Generator(device=_env.DEVICE).manual_seed(GLOBAL_SEED) + cell = torch.eye(3, dtype=torch.float64, device=_env.DEVICE) * 6.0 + coord = torch.rand( + [6, 3], dtype=torch.float64, device=_env.DEVICE, generator=generator + ) + coord = torch.matmul(coord, cell).unsqueeze(0) + atype = torch.tensor( + [[0, 0, 0, 1, 1, 1]], dtype=torch.int64, device=_env.DEVICE + ) + spin = 0.1 * torch.rand( + [1, 6, 3], dtype=torch.float64, device=_env.DEVICE, generator=generator + ) + box = cell.reshape(1, 9) + + e_excl = excluded(coord, atype, spin, box=box)["energy"] + e_base = baseline(coord, atype, spin, box=box)["energy"] + de = (e_excl - e_base).abs().max().item() + assert de > 1e-6, ( + f"pair_exclude_types=[[0, 1]] left the energy unchanged (diff=" + f"{de:.3e}); the exclusion never reached the neighbor-graph build" + ) diff --git a/source/tests/pt_expt/model/test_get_model_dpa4.py b/source/tests/pt_expt/model/test_get_model_dpa4.py index 210cec431f..aa76fd7ecc 100644 --- a/source/tests/pt_expt/model/test_get_model_dpa4.py +++ b/source/tests/pt_expt/model/test_get_model_dpa4.py @@ -192,10 +192,13 @@ def test_pair_exclude_types_mismatch_raises(self) -> None: get_model(raw) def test_unsupported_keys_raise(self) -> None: - """pt-only SeZM model-level features fail fast with NotImplementedError.""" + """pt-only SeZM model-level features fail fast with NotImplementedError. + + ``bridging_method`` is no longer in this list: it is supported as an + atomic-model composition (see ``test_zbl_bridging.py``). + """ cases = { "spin": ({"use_spin": [True, False], "virtual_scale": [0.3]}, "Spin DPA4"), - "bridging_method": ("ZBL", "`bridging_method` is not supported"), "lora": ({"rank": 4}, "`lora` is not supported"), "use_compile": (True, "`use_compile` is not supported"), "preset_out_bias": ( @@ -209,6 +212,54 @@ def test_unsupported_keys_raise(self) -> None: with self.assertRaisesRegex(NotImplementedError, msg_regex): get_model(raw) + def test_native_spin_capability_gate_standard_config(self) -> None: + """The generic ``supports_native_spin()`` gate rejects a dense descriptor. + + Uses a complete ``type="standard"`` se_e2_a energy config so the + DESCRIPTOR-AGNOSTIC capability gate is what fires -- not the + dpa4-typed builder's descriptor contract (pinned separately below). + """ + raw = { + "type": "standard", + "type_map": ["Ni", "O"], + "descriptor": { + "type": "se_e2_a", + "rcut": 4.0, + "rcut_smth": 3.5, + "sel": [8, 8], + }, + "fitting_net": {"type": "ener", "neuron": [8, 8]}, + "spin": {"use_spin": [True, False], "scheme": "native"}, + } + with self.assertRaisesRegex(NotImplementedError, "native spin"): + get_model(raw) + + def test_native_spin_non_dpa4_descriptor_raises(self) -> None: + """A dpa4-typed config rejects a foreign descriptor (family contract). + + The dpa4-typed builder pins its descriptor/fitting contract before + the generic capability gate is reached, so the mismatch surfaces as + the family builder's ``ValueError``. + """ + raw = _make_raw_model_config() + raw["descriptor"] = {"type": "se_e2_a"} + raw["spin"] = {"use_spin": [True, False], "scheme": "native"} + with self.assertRaisesRegex(ValueError, "DPA4/SeZM descriptor"): + get_model(raw) + + def test_native_spin_add_chg_spin_ebd_combined_builds(self) -> None: + """Native-scheme spin combined with charge-spin FiLM is SUPPORTED. + + (Review 3638047227 lifted the old rejection; the combined model's + behavior is pinned in ``test_dpa4_native_spin.py``.) + """ + raw = _make_raw_model_config() + raw["descriptor"]["add_chg_spin_ebd"] = True + raw["spin"] = {"use_spin": [True, False], "scheme": "native"} + model = get_model(raw) + self.assertTrue(model.has_chg_spin_ebd()) + self.assertTrue(model.has_spin()) + def test_default_unsupported_values_pass(self) -> None: """Normalized defaults (bridging None, lora None, use_compile False) build.""" model_params = _normalize_model(_make_raw_model_config()) @@ -223,7 +274,7 @@ def test_default_unsupported_values_pass(self) -> None: # `enable_tf32` toggles TF32 matmul precision in pt but is ignored by pt_expt # (always "highest" precision); a truthy value must emit a warn-once message. @pytest.mark.parametrize("enable_tf32", [True, False]) # truthy warns, falsy silent -def test_enable_tf32_warns_once(enable_tf32, caplog, monkeypatch) -> None: +def test_enable_tf32_warns_once(enable_tf32, monkeypatch) -> None: import importlib # the package __init__ rebinds the name ``get_model`` to the function, so @@ -234,20 +285,50 @@ def test_enable_tf32_warns_once(enable_tf32, caplog, monkeypatch) -> None: # test ordering (other get_sezm_model calls may have already warned) monkeypatch.setattr(gm_mod, "_WARNED_ONCE", set()) - raw = _make_raw_model_config(enable_tf32=enable_tf32) + # Count emissions on the EMITTING logger with our own handler rather than + # through caplog: caplog reads a root handler, so whatever global logging + # state earlier tests left behind (set_log_handles flips the ``deepmd`` + # logger's propagate off and installs its own handlers) changes how many + # records reach it -- zero when propagation is off, more than one when the + # record is seen through several attached handlers. A handler on the + # emitting logger sees exactly one record per ``log.warning`` call. + records: list[logging.LogRecord] = [] + + class _Collect(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + records.append(record) - with caplog.at_level(logging.WARNING, logger=gm_mod.log.name): - gm_mod.get_sezm_model(raw) - matches = [r for r in caplog.records if "enable_tf32" in r.getMessage()] - if enable_tf32: - assert len(matches) == 1, caplog.text - # a second call must NOT warn again (warn-once per process) - caplog.clear() - with caplog.at_level(logging.WARNING, logger=gm_mod.log.name): + handler = _Collect(level=logging.WARNING) + old_level = gm_mod.log.level + gm_mod.log.setLevel(logging.WARNING) + gm_mod.log.addHandler(handler) + try: + gm_mod.get_sezm_model(_make_raw_model_config(enable_tf32=enable_tf32)) + matches = [r for r in records if "enable_tf32" in r.getMessage()] + if enable_tf32: + assert len(matches) == 1, [r.getMessage() for r in records] + # a second call must NOT warn again (warn-once per process) + records.clear() gm_mod.get_sezm_model(_make_raw_model_config(enable_tf32=enable_tf32)) - assert not [r for r in caplog.records if "enable_tf32" in r.getMessage()] - else: - assert not matches, caplog.text + assert not [r for r in records if "enable_tf32" in r.getMessage()] + else: + assert not matches, [r.getMessage() for r in records] + finally: + gm_mod.log.removeHandler(handler) + gm_mod.log.setLevel(old_level) + + +class TestNativeSpinErrorTranslation(unittest.TestCase): + """Only the unexpected-``use_spin`` TypeError becomes the capability error.""" + + def test_unrelated_construction_error_propagates(self) -> None: + # A bogus fitting kwarg must surface as the REAL TypeError, not be + # masked as a native-spin capability failure (review 3644847676). + raw = _make_raw_model_config() + raw["spin"] = {"use_spin": [True, False], "scheme": "native"} + raw["fitting_net"]["bogus_option"] = 1 + with self.assertRaisesRegex(TypeError, "bogus_option"): + get_model(raw) if __name__ == "__main__": diff --git a/source/tests/pt_expt/model/test_graph_export_with_comm.py b/source/tests/pt_expt/model/test_graph_export_with_comm.py index d1c031074c..aa73a3f46e 100644 --- a/source/tests/pt_expt/model/test_graph_export_with_comm.py +++ b/source/tests/pt_expt/model/test_graph_export_with_comm.py @@ -16,6 +16,7 @@ import pytest +from deepmd.pt_expt.model.get_model import get_model as _get_pt_expt_model from deepmd.pt_expt.utils.env import ( DEVICE, ) @@ -23,6 +24,10 @@ deserialize_to_file, ) +from .test_dpa4_export import ( + _DPA4_CONFIG, +) + # Small graph-eligible dpa2 descriptor: tebd_input_mode defaults to # "concat", use_three_body defaults to False -> uses_graph_lower() is True, # and has_message_passing_across_ranks() is unconditionally True for any @@ -135,6 +140,49 @@ def test_dpa2_graph_pt2_embeds_with_comm_artifact(dpa2_dpmodel_data, tmp_path) - assert meta["has_message_passing"] is True +@pytest.fixture(scope="module") +def dpa4_pt_expt_data() -> dict: + """Build a serialized pt_expt DPA4/SeZM model dict (same shape as + ``dp freeze`` input) -- mirrors ``test_dpa4_export.py``'s construction. + + DPA4/SeZM's ``"type": "dpa4"`` model key is only special-cased by the + pt_expt model factory (``get_sezm_model``), not the plain dpmodel one + ``_build_data`` above uses for dpa1/dpa2, so this fixture builds the + model directly via ``deepmd.pt_expt.model.get_model.get_model`` + instead. + """ + model = _get_pt_expt_model(copy.deepcopy(_DPA4_CONFIG)) + model.to("cpu") + model.eval() + return {"model": model.serialize()} + + +def test_dpa4_graph_pt2_embeds_with_comm_artifact(dpa4_pt_expt_data, tmp_path) -> None: + # DPA4 (graph-only comm family): the graph freeze embeds the nested + # with-comm artifact; regression companion to the dpa1 no-artifact test. + # + # DPA4's cross-rank ghost exchange is implemented ONLY on the graph + # lower (``has_message_passing_across_ranks()`` is True, but + # ``dense_lower_supports_comm()`` is False -- see the dpa4 descriptor + # and ``test_dpa4_export.py``'s module docstring); freezing with + # ``lower_kind="graph"`` must therefore embed the nested with-comm + # artifact, unlike the dpa1 (non-message-passing) graph freeze above. + p = str(tmp_path / "m_dpa4_graph.pt2") + deserialize_to_file( + p, + copy.deepcopy(dpa4_pt_expt_data), + lower_kind="graph", + ) + with zipfile.ZipFile(p, "r") as zf: + names = zf.namelist() + assert "model/extra/forward_lower_with_comm.pt2" in names + + meta = _read_metadata(p) + assert meta["lower_input_kind"] == "graph" + assert meta["has_comm_artifact"] is True + assert meta["has_message_passing"] is True + + def test_dpa2_graph_with_comm_aparam_freeze(tmp_path) -> None: """A ``numb_aparam > 0`` message-passing model freezes to the graph ``.pt2`` with the nested with-comm artifact. @@ -310,3 +358,96 @@ def run(n_local_val: int) -> torch.Tensor: # keepalive: the raw pointer in ``comm`` must reference a live buffer # through both ``run`` calls above (a real use, not a bare ``del``) assert sendlist_indices.ctypes.data_as(ctypes.c_void_p).value == addr + + +# --- native spin: multi-rank on the graph lower -------------------------- +# pt's ``SeZMModel.supports_edge_parallel`` is NOT overridden by +# ``SeZMNativeSpinModel``, so native spin participates in the with-comm +# artifact there; pt_expt used to exclude every ``NativeSpinModelKind`` +# unconditionally, which silently made native spin single-rank only. The +# spin input is per-node and its ghost rows arrive via the LAMMPS ``sp`` +# forward-comm, so spin itself needs no cross-rank exchange -- the per-block +# ghost FEATURE refresh is the same ``border_op`` the energy model drives. +_NATIVE_SPIN_CONFIG = { + **copy.deepcopy(_DPA4_CONFIG), + "spin": {"use_spin": [True, False], "scheme": "native"}, +} + + +def _native_spin_model(): + model = _get_pt_expt_model(copy.deepcopy(_NATIVE_SPIN_CONFIG)) + return model.to("cpu").eval() + + +def test_native_spin_needs_with_comm_on_the_graph_lower_only() -> None: + """The gate admits native spin on graph, still refuses it on nlist. + + Native spin has no dense with-comm wrapper at all, so the nlist branch + must stay ``False`` -- otherwise the freeze would try to trace a lower + that does not exist. + """ + from deepmd.pt_expt.utils.serialization import ( + _needs_with_comm_artifact, + ) + + model = _native_spin_model() + assert _needs_with_comm_artifact(model, "graph") is True + assert _needs_with_comm_artifact(model, "nlist") is False + + +def test_native_spin_graph_with_comm_abi() -> None: + """The with-comm graph-spin ABI is the non-comm one plus the comm block. + + 22 positional inputs: the 14-slot spin graph base (10 CSR + ``spin`` at + slot 10 + the None-valued fparam/aparam/charge_spin tail) followed by + the 8 comm tensors -- i.e. ``spin`` keeps slot 10 and the comm block + starts at 14. Pinning the count is what keeps the C++ feeder and this + trace from drifting apart. + """ + from deepmd.dpmodel.model.model import get_model as _get_dp_model + from deepmd.pt_expt.utils.serialization import ( + _trace_and_export, + ) + + dp_cfg = copy.deepcopy(_NATIVE_SPIN_CONFIG) + dp_cfg.pop("type", None) # generic builder; the dpa4 alias routes alike + data = { + "model": _get_dp_model(dp_cfg).serialize(), + "model_def_script": copy.deepcopy(_NATIVE_SPIN_CONFIG), + } + exported, _meta, _dj, keys = _trace_and_export( + data, + model_json_override=None, + with_comm_dict=True, + lower_kind="graph", + ) + placeholders = exported.module().graph.find_nodes(op="placeholder") + assert len(placeholders) == 22, ( + f"graph-spin with-comm program must accept 22 positional inputs " + f"(14 spin-graph base incl. spin at slot 10 + 8 comm); got " + f"{len(placeholders)}" + ) + # the spin-specific outputs must survive the with-comm trace + for key in ("energy", "force", "force_mag", "mask_mag"): + assert key in keys, f"{key} missing from the with-comm graph outputs" + + +def test_graph_eligibility_guard_is_stated_exactly_once() -> None: + """The innermost ``lower_kind="graph"`` guard must not be duplicated. + + It was once present as two verbatim copies from separate commits; the + second could never fire yet invited the two to drift. Asserted on the + source text -- behavior cannot distinguish one guard from two. Read from + the file, not ``inspect.getsource``: conftest wraps ``_trace_and_export``. + """ + from pathlib import ( + Path, + ) + + from deepmd.pt_expt.utils import ( + serialization, + ) + + source = Path(serialization.__file__).read_text() + assert source.count("graph-lower eligible (model_uses_graph_lower() is False") == 1 + assert source.count("if not model_uses_graph_lower(model):") == 1 diff --git a/source/tests/pt_expt/model/test_zbl_bridging.py b/source/tests/pt_expt/model/test_zbl_bridging.py new file mode 100644 index 0000000000..8151c88f7f --- /dev/null +++ b/source/tests/pt_expt/model/test_zbl_bridging.py @@ -0,0 +1,760 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""pt_expt ZBL bridging as COMPOSITION (review 3638077323, redesigned). + +``bridging_method: ZBL`` builds a linear composition +(``LinearEnergyModel`` over ``[learned, InterPotentialAtomicModel]`` with +``weights="sum"``); eager values still match pt's flag-architected +``SeZMModel`` bit-for-bit (identical math), pinned here as a value +regression together with FD force, export/DeepEval e2e, training smoke, +and the single-rank with-comm gate. +""" + +import copy +import json +import os +import zipfile + +import numpy as np +import pytest +import torch + +from deepmd.pt.model.model import get_model as pt_get_model +from deepmd.pt_expt.descriptor.dpa4 import ( + DescrptDPA4, +) +from deepmd.pt_expt.fitting.dpa4_ener import ( + SeZMEnergyFittingNet, +) +from deepmd.pt_expt.model.dp_linear_model import ( + LinearEnergyModel, +) +from deepmd.pt_expt.model.get_model import ( + get_model, +) + +from ...seed import ( + GLOBAL_SEED, +) + +ZBL_CONFIG = { + "type": "dpa4", + "type_map": ["Ni", "O"], + "descriptor": { + "type": "dpa4", + "rcut": 4.0, + "sel": 8, + "channels": 16, + "n_radial": 8, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "precision": "float64", + "seed": 7, + "random_gamma": False, + }, + "fitting_net": { + "type": "dpa4_ener", + "neuron": [8, 8], + "precision": "float64", + "seed": 7, + }, + "bridging_method": "ZBL", + "bridging_r_inner": 0.8, + "bridging_r_outer": 1.2, +} + + +def _analytic_zbl_total(coord, atype, rcut, type_map=("Ni", "O")) -> float: + """Independent in-test ZBL reference: direct double loop over pairs.""" + import math + + z_of = {"Ni": 28.0, "O": 8.0} + zs = [z_of[type_map[t]] for t in atype] + a_coeff = (0.18175, 0.50986, 0.28022, 0.028171) + b_coeff = (3.1998, 0.94229, 0.4029, 0.20162) + total = 0.0 + n = len(atype) + for i in range(n): + for j in range(i + 1, n): + r = float(np.linalg.norm(coord[i] - coord[j])) + if r >= rcut: + continue + a = 0.88534 * 0.5291772109 / (zs[i] ** 0.23 + zs[j] ** 0.23) + phi = sum( + ak * math.exp(-bk * (r / a)) + for ak, bk in zip(a_coeff, b_coeff, strict=True) + ) + total += 14.3996 * zs[i] * zs[j] / r * phi + return total + + +def _close_pair_system(cpu): + generator = torch.Generator(device=cpu).manual_seed(GLOBAL_SEED + 2) + nloc = 6 + cell = torch.rand([3, 3], dtype=torch.float64, generator=generator) + cell = (cell + cell.T) + 6.0 * torch.eye(3) + coord = 1.5 + 3.0 * torch.rand([nloc, 3], dtype=torch.float64, generator=generator) + coord[1] = coord[0] + torch.tensor([0.95, 0.0, 0.0], dtype=torch.float64) + atype = torch.tensor([[0, 0, 1, 0, 1, 1]], dtype=torch.int64) + return coord.unsqueeze(0), atype, cell.reshape(1, 9) + + +class TestZBLBridgingPtExpt: + def setup_method(self) -> None: + cpu = torch.device("cpu") + pt_model = pt_get_model(copy.deepcopy(ZBL_CONFIG)).to(torch.float64) + # JITTER the reference weights: a fresh DPA4 zero-initializes its + # residual projections and is architecturally input-independent in + # those paths, which would make the parity below partially vacuous + # (see dpa4_fixtures.jitter_zero_arrays). + from deepmd.pt.model.descriptor.sezm import ( + DescrptSeZM, + ) + + from ...dpa4_fixtures import ( + jitter_zero_arrays, + ) + + jittered = jitter_zero_arrays( + pt_model.atomic_model.descriptor.serialize(), np.random.default_rng(3) + ) + pt_model.atomic_model.descriptor = DescrptSeZM.deserialize(jittered).to( + torch.float64 + ) + self.pt_model = pt_model.eval().to(cpu) + assert self.pt_model.inter_potential is not None + + pt_expt_model = get_model(copy.deepcopy(ZBL_CONFIG)) + assert type(pt_expt_model) is LinearEnergyModel + dp_child = pt_expt_model.atomic_model.models[0] + # weight copy into the LEARNED child: pt DescrptSeZM / fitting + # serialize to the SAME backend-agnostic dict schema (incl. the + # InnerClamp radii) + dp_child.descriptor = DescrptDPA4.deserialize( + self.pt_model.atomic_model.descriptor.serialize() + ) + dp_child.fitting_net = SeZMEnergyFittingNet.deserialize( + self.pt_model.atomic_model.fitting_net.serialize() + ) + self.pt_expt_model = pt_expt_model.to(cpu).eval() + self.coord, self.atype, self.box = _close_pair_system(cpu) + + def test_parity_vs_pt_with_zbl(self) -> None: + """Composition == pt's flag architecture on the same weights (values). + + pt adds the raw ZBL to the fitting energy; the composition sums the + same two per-atom terms -- identical math, pinned at 1e-12 for + energy/force/virial. + """ + out_pt = self.pt_model.forward(self.coord, self.atype, self.box) + out_pte = self.pt_expt_model.forward(self.coord, self.atype, box=self.box) + # anti-vacuity: the jittered network must produce nontrivial forces, + # else the parity would compare zeros with zeros. + assert out_pte["force"].abs().max().item() > 1e-6 + for key in ("energy", "force", "virial"): + torch.testing.assert_close( + out_pt[key], out_pte[key], rtol=1e-12, atol=1e-12, msg=key + ) + + def test_zbl_child_adds_positive_energy(self) -> None: + """Learned child alone vs the composition: positive ZBL repulsion.""" + from deepmd.pt_expt.model.ener_model import ( + EnergyModel, + ) + + m_dp = ( + EnergyModel(atomic_model_=self.pt_expt_model.atomic_model.models[0]) + .to(torch.device("cpu")) + .eval() + ) + e_sum = self.pt_expt_model.forward(self.coord, self.atype, box=self.box)[ + "energy" + ] + e_dp = m_dp.forward(self.coord, self.atype, box=self.box)["energy"] + diff = float((e_sum - e_dp).sum()) + # EXACT analytical check, not just positivity: the composition's + # extra term must equal the independently computed ZBL sum over all + # pairs within rcut (gas phase: no box, so a direct double loop is + # the complete reference). + e_gas_sum = self.pt_expt_model.forward(self.coord, self.atype)["energy"] + e_gas_dp = m_dp.forward(self.coord, self.atype)["energy"] + ref = _analytic_zbl_total( + self.coord[0].numpy(), self.atype[0].numpy(), rcut=4.0 + ) + np.testing.assert_allclose( + float((e_gas_sum - e_gas_dp).sum()), ref, rtol=1e-10, atol=1e-10 + ) + assert diff > 1e-3 + + def test_force_matches_finite_difference(self) -> None: + """F = -dE/dx through the shared-edge-leaf summed autograd.""" + eps = 1e-5 + out = self.pt_expt_model.forward(self.coord, self.atype, box=self.box) + force = out["force"].reshape(-1, 3) + for atom, comp in ((1, 0), (2, 2)): # close-pair atom + a far atom + cp = self.coord.clone() + cp[0, atom, comp] += eps + ep = self.pt_expt_model.forward(cp, self.atype, box=self.box)["energy"] + cm = self.coord.clone() + cm[0, atom, comp] -= eps + em = self.pt_expt_model.forward(cm, self.atype, box=self.box)["energy"] + fd = -float((ep - em).sum()) / (2 * eps) + np.testing.assert_allclose( + float(force[atom, comp]), fd, rtol=1e-6, atol=1e-6 + ) + + def test_serialize_roundtrip(self) -> None: + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + + data = self.pt_expt_model.serialize() + # the flat wire type is "linear" -- the SAME string pt/tf write, so a + # composition round-trips across backends + assert data["type"] == "linear" + m2 = BaseModel.deserialize(data).to(torch.device("cpu")).eval() + assert type(m2) is LinearEnergyModel + out = self.pt_expt_model.forward(self.coord, self.atype, box=self.box) + out2 = m2.forward(self.coord, self.atype, box=self.box) + torch.testing.assert_close( + out["energy"], out2["energy"], rtol=1e-12, atol=1e-12 + ) + + def test_with_comm_gate_off_for_composition(self) -> None: + """Compositions never compile a with-comm artifact (single-rank).""" + from deepmd.pt_expt.utils.serialization import ( + _needs_with_comm_artifact, + ) + + assert ( + _needs_with_comm_artifact(self.pt_expt_model, lower_kind="graph") is False + ) + + def test_pt_bridging_checkpoint_rejected(self) -> None: + """Reject pt's flag-serialized bridging checkpoints. + + pt serializes bridging as a wrapper flag; our architecture is a + linear composition with a different dict shape -- fail fast instead + of a silent wrong conversion. + """ + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + + with pytest.raises(NotImplementedError, match="bridging_method"): + BaseModel.deserialize(self.pt_model.serialize()) + + +def _native_spin_zbl_config() -> dict: + cfg = copy.deepcopy(ZBL_CONFIG) + cfg["spin"] = {"use_spin": [True, False], "scheme": "native"} + return cfg + + +def _spin_system(): + """6 atoms (3 Ni spin-active, 3 O) with one close pair driving the ZBL.""" + coord = torch.tensor( + [ + [ + [1.0, 1.0, 1.0], + [1.9, 1.2, 1.1], # close to atom 0 -> nontrivial ZBL + [3.0, 2.0, 1.0], + [1.0, 3.0, 2.0], + [3.5, 1.0, 2.0], + [2.0, 2.0, 3.0], + ] + ], + dtype=torch.float64, + ) + atype = torch.tensor([[0, 0, 0, 1, 1, 1]], dtype=torch.int64) + spin = 0.1 * torch.ones_like(coord) + box = (6.0 * torch.eye(3, dtype=torch.float64)).reshape(1, 9) + return coord, atype, spin, box + + +class TestNativeSpinWithBridging: + """Native spin + analytical bridging compose (review 3649276109). + + ``get_standard_model`` OWNS assembling the atomic model, bridging + composition included, and the native-spin wrapper re-classes whatever it + returns -- so the two features combine with no special case: the learned + child consumes ``spin``, the analytical child accepts and ignores it. + """ + + def test_construction_composes_and_keeps_spin(self) -> None: + from deepmd.pt_expt.model.native_spin_model import ( + NativeSpinEnergyModel, + ) + + model = get_model(_native_spin_zbl_config()) + assert isinstance(model, NativeSpinEnergyModel) + assert model.has_spin() is True + kinds = [type(c).__name__ for c in model.atomic_model.models] + assert kinds[1] == "InterPotentialAtomicModel", kinds + # bridging radii still reach the LEARNED child's descriptor + assert float(model.atomic_model.models[0].descriptor.inner_clamp.r_inner) == 0.8 + + def test_forward_energy_force_force_mag(self) -> None: + model = get_model(_native_spin_zbl_config()).to(torch.device("cpu")).eval() + coord, atype, spin, box = _spin_system() + out = model(coord, atype, spin, box=box) + for key in ("energy", "force", "force_mag"): + assert torch.isfinite(out[key]).all(), key + # mask_mag follows use_spin=[True, False] on atype [0,0,0,1,1,1] + assert out["mask_mag"].reshape(-1).tolist() == [ + True, + True, + True, + False, + False, + False, + ] + # anti-vacuity: the analytical term must actually contribute. The + # learned child alone is the same model minus the ZBL energy. + from deepmd.pt_expt.model.native_spin_model import ( + NativeSpinEnergyModel, + ) + + learned_only = ( + NativeSpinEnergyModel( + atomic_model_=model.atomic_model.models[0], spin=model.spin + ) + .to(torch.device("cpu")) + .eval() + ) + # Gas phase (no box) for the EXACT check: the in-test reference is a + # direct double loop over pairs, which has no periodic images. + e_gas = model(coord, atype, spin)["energy"] + e_gas_learned = learned_only(coord, atype, spin)["energy"] + zbl_contrib = float((e_gas - e_gas_learned).sum()) + ref = _analytic_zbl_total(coord[0].numpy(), atype[0].numpy(), rcut=4.0) + np.testing.assert_allclose(zbl_contrib, ref, rtol=1e-10, atol=1e-10) + assert zbl_contrib > 1e-3 + + def test_serialize_roundtrip(self) -> None: + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + + model = get_model(_native_spin_zbl_config()).to(torch.device("cpu")).eval() + data = model.serialize() + assert data["type"] == "native_spin" + restored = BaseModel.deserialize(data).to(torch.device("cpu")).eval() + coord, atype, spin, box = _spin_system() + out = model(coord, atype, spin, box=box) + out2 = restored(coord, atype, spin, box=box) + for key in ("energy", "force", "force_mag"): + torch.testing.assert_close( + out[key], out2[key], rtol=1e-12, atol=1e-12, msg=key + ) + + +def test_native_spin_with_bridging_dpmodel() -> None: + """Dpmodel twin: same composition, energy-only (no autograd there).""" + from deepmd.dpmodel.model.model import get_model as dp_get_model + from deepmd.dpmodel.model.native_spin_model import ( + NativeSpinEnergyModel as NativeSpinEnergyModelDP, + ) + + cfg = _native_spin_zbl_config() + cfg.pop("type") # generic builder; the dpa4 alias routes the same way + model = dp_get_model(cfg) + assert isinstance(model, NativeSpinEnergyModelDP) + kinds = [type(c).__name__ for c in model.atomic_model.models] + assert kinds[1] == "InterPotentialAtomicModel", kinds + + coord, atype, spin, box = _spin_system() + out = model.call(coord.numpy(), atype.numpy(), spin.numpy(), box=box.numpy()) + assert np.all(np.isfinite(out["energy"])) + assert out["mask_mag"].reshape(-1).tolist() == [ + True, + True, + True, + False, + False, + False, + ] + + +def test_bridging_radii_defaults() -> None: + """bridging_r_inner/r_outer default to 0.5/0.8 on the learned child.""" + cfg = copy.deepcopy(ZBL_CONFIG) + cfg.pop("bridging_r_inner") + cfg.pop("bridging_r_outer") + model = get_model(cfg) + ic = model.atomic_model.models[0].descriptor.inner_clamp + assert ic is not None + assert float(ic.r_inner) == 0.5 + assert float(ic.r_outer) == 0.8 + + +class TestZBLBridgingExportAndTraining: + """Graph .pt2 freeze + DeepEval parity and a trainer smoke.""" + + def test_graph_freeze_and_deep_eval_parity(self, tmp_path) -> None: + if os.environ.get("CI") == "true": + pytest.skip( + "AOTInductor compile is slow (minutes); local/fixture-gen only." + ) + from deepmd.infer import ( + DeepPot, + ) + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, + ) + + cpu = torch.device("cpu") + model = get_model(copy.deepcopy(ZBL_CONFIG)).to(cpu).eval() + coord, atype, box = _close_pair_system(cpu) + ref = model.forward(coord, atype, box=box) + + model_file = tmp_path / "dpa4_zbl_graph.pt2" + data = {"model": model.serialize()} + deserialize_to_file(str(model_file), data, lower_kind="graph") + + with zipfile.ZipFile(model_file) as z: + md = json.loads(z.read("model/extra/metadata.json").decode("utf-8")) + # single-rank contract: compositions never get a with-comm artifact + assert md["has_comm_artifact"] is False + + dp = DeepPot(str(model_file)) + e, f, v = dp.eval( + coord.reshape(1, -1).numpy(), + box.numpy(), + atype.reshape(-1).numpy(), + atomic=False, + ) + np.testing.assert_allclose( + np.asarray(e).reshape(-1), + ref["energy"].detach().numpy().reshape(-1), + rtol=1e-10, + atol=1e-10, + err_msg="energy", + ) + np.testing.assert_allclose( + np.asarray(f).reshape(-1), + ref["force"].detach().numpy().reshape(-1), + rtol=1e-10, + atol=1e-10, + err_msg="force", + ) + + def test_training_smoke(self, tmp_path) -> None: + data_dir = os.path.join( + os.path.dirname(__file__), "..", "..", "pt", "NiO", "data", "single" + ) + if not os.path.isdir(data_dir): + pytest.skip(f"NiO data not found: {data_dir}") + from deepmd.pt_expt.entrypoints.main import ( + get_trainer, + ) + from deepmd.pt_expt.train.training import ( + DEFAULT_TASK_KEY, + ) + from deepmd.utils.argcheck import ( + normalize, + ) + from deepmd.utils.compat import ( + update_deepmd_input, + ) + + config = { + "model": copy.deepcopy(ZBL_CONFIG), + "learning_rate": { + "type": "exp", + "decay_steps": 500, + "start_lr": 0.001, + "stop_lr": 3.51e-8, + }, + "loss": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + }, + "training": { + "training_data": {"systems": [data_dir], "batch_size": 1}, + "validation_data": { + "systems": [data_dir], + "batch_size": 1, + "numb_btch": 1, + }, + "numb_steps": 2, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 1, + "save_freq": 2, + }, + } + config = update_deepmd_input(config, warning=False) + config = normalize(config) + + old_cwd = os.getcwd() + os.chdir(tmp_path) + try: + trainer = get_trainer(config) + model = trainer.wrapper.model[DEFAULT_TASK_KEY] + assert type(model) is LinearEnergyModel + tasks = trainer._make_training_tasks() + task = trainer.select_task(tasks) + for step in range(2): + result = trainer.train_step(task, step) + loss = result.payload["loss"] + assert torch.isfinite(loss).all(), f"non-finite loss at step {step}" + finally: + os.chdir(old_cwd) + + +class TestInterPotentialChangeTypeMapPtExpt: + """pt_expt twin of the dpmodel ``change_type_map`` regression. + + Exercised through the REAL composition: ``LinearEnergyModel`` -> + ``LinearEnergyAtomicModel`` -> ``InterPotentialAtomicModel`` -> + ``InterPotential``. Inside a pt_expt module tree the element lookup is a + wrapped torch buffer, so the rebuild must land on the same + device/namespace (review 3649295675) -- a numpy rebuild would desync the + buffer or fail outright on CUDA. + """ + + @staticmethod + def _zbl_child(model): + return model.atomic_model.models[1] + + @staticmethod + def _pair_energy(zbl_child, atype_value: int = 0, r: float = 1.0) -> float: + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + ) + from deepmd.pt_expt.utils import env as _env + + graph = NeighborGraph( + n_node=torch.tensor([2], dtype=torch.int64, device=_env.DEVICE), + edge_index=torch.tensor( + [[0, 1], [1, 0]], dtype=torch.int64, device=_env.DEVICE + ), + edge_vec=torch.tensor( + [[r, 0.0, 0.0], [-r, 0.0, 0.0]], + dtype=torch.float64, + device=_env.DEVICE, + ), + edge_mask=torch.ones(2, dtype=torch.bool, device=_env.DEVICE), + ) + atype = torch.full((2,), atype_value, dtype=torch.int64, device=_env.DEVICE) + return float( + zbl_child.forward_common_atomic_graph(graph, atype)["energy"].sum() + ) + + def _build(self, type_map): + from deepmd.pt_expt.utils import env as _env + + config = copy.deepcopy(ZBL_CONFIG) + config["type_map"] = list(type_map) + return get_model(config).to(_env.DEVICE).eval() + + def test_lookup_is_a_wrapped_buffer(self) -> None: + """Precondition: inside pt_expt the lookup is a torch buffer.""" + from deepmd.pt_expt.utils import env as _env + + z = self._zbl_child(self._build(["Ni", "O"])).potential.atomic_numbers + assert isinstance(z, torch.Tensor), ( + "the pt_expt wrapper no longer converts the lookup to a tensor; " + "this test would stop covering the device-safe rebuild" + ) + assert z.device.type == torch.device(_env.DEVICE).type + + def test_reorder_matches_a_freshly_built_model(self) -> None: + # NOTE: applied to the ZBL CHILD, not the whole composition -- the + # DPA4/SeZM learned child does not implement change_type_map at all + # ("change_type_map is not supported for SeZM"), a separate pre-existing + # limitation. The lookup under test belongs to this child. + child = self._zbl_child(self._build(["Ni", "O"])) + e_nini = self._pair_energy(child) + child.change_type_map(["O", "Ni"]) + fresh = self._zbl_child(self._build(["O", "Ni"])) + e_fresh = self._pair_energy(fresh) + # anti-vacuity: Ni-Ni and O-O must be far apart, else a stale lookup + # would be indistinguishable from a rebuilt one + assert abs(e_fresh - e_nini) > 1.0 + np.testing.assert_allclose(self._pair_energy(child), e_fresh, rtol=1e-12) + assert [float(v) for v in child.potential.atomic_numbers] == [8.0, 28.0] + + def test_added_element_extends_the_lookup_on_device(self) -> None: + from deepmd.pt_expt.utils import env as _env + + child = self._zbl_child(self._build(["Ni", "O"])) + child.change_type_map(["Ni", "O", "H"]) + z = child.potential.atomic_numbers + assert isinstance(z, torch.Tensor), "the rebuild dropped out of torch" + assert z.device.type == torch.device(_env.DEVICE).type + assert [float(v) for v in z] == [28.0, 8.0, 1.0] + # the new type is addressable -- a stale (length-2) table raises here + np.testing.assert_allclose( + self._pair_energy(child, atype_value=2), + self._pair_energy(self._zbl_child(self._build(["H"]))), + rtol=1e-12, + ) + + def test_serialize_roundtrip_after_change_type_map(self) -> None: + """Checkpoint continuity: the restored child must predict the same. + + Serialization records the NEW public map, so a stale in-memory lookup + and its deserialized twin disagree -- the restart-time symptom. + """ + from deepmd.dpmodel.atomic_model.base_atomic_model import ( + BaseAtomicModel, + ) + + child = self._zbl_child(self._build(["Ni", "O"])) + child.change_type_map(["O", "Ni"]) + data = child.serialize() + # The wire type must still resolve through the shared registry ... + assert BaseAtomicModel.get_class_by_type(data["type"]) is not None + # ... but restore through the pt_expt class the child actually is: + # the dpmodel class is NumPy-backed, and this test feeds device + # tensors, so a CUDA run would index a NumPy out_bias with a CUDA + # atype and fail. Same-class restore is also the stricter check. + restored = type(child).deserialize(data) + assert restored.get_type_map() == ["O", "Ni"] + np.testing.assert_allclose( + self._pair_energy(restored), self._pair_energy(child), rtol=1e-12 + ) + + +def test_native_spin_with_bridging_graph_freeze_and_deep_eval(tmp_path) -> None: + """Native spin + ZBL freezes to a graph .pt2 and evaluates in parity. + + Both features are graph-route-only, so their combination must survive the + export seam too, not just eager construction (review 3649276109). + """ + if os.environ.get("CI") == "true": + pytest.skip("AOTInductor compile is slow (minutes); local/fixture-gen only.") + from deepmd.infer import ( + DeepPot, + ) + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, + ) + + cpu = torch.device("cpu") + model = get_model(_native_spin_zbl_config()).to(cpu).eval() + coord, atype, spin, box = _spin_system() + ref = model(coord, atype, spin, box=box) + + model_file = tmp_path / "dpa4_native_spin_zbl_graph.pt2" + # native spin has no dense lower at all, so the graph kind is the only + # valid one here; the composition additionally forbids a with-comm twin. + deserialize_to_file( + str(model_file), {"model": model.serialize()}, lower_kind="graph" + ) + with zipfile.ZipFile(model_file) as z: + md = json.loads(z.read("model/extra/metadata.json").decode("utf-8")) + assert md["is_spin"] is True + assert md["has_comm_artifact"] is False + assert md["use_spin"] == [True, False] + + dp = DeepPot(str(model_file)) + assert dp.has_spin + e, f, v, fm, mm = dp.eval( + coord.numpy(), + box.numpy(), + atype.reshape(-1).numpy(), + atomic=False, + spin=spin.numpy(), + )[:5] + for got, want, name in ( + (e, ref["energy"], "energy"), + (f, ref["force"], "force"), + (fm, ref["force_mag"], "force_mag"), + ): + np.testing.assert_allclose( + np.asarray(got).reshape(-1), + want.detach().numpy().reshape(-1), + rtol=1e-10, + atol=1e-10, + err_msg=name, + ) + + +def test_bridged_metadata_carries_charge_spin_dim(tmp_path) -> None: + """The FROZEN metadata must declare charge_spin for a bridged model. + + This is the consequence the eager forward hides: the learned child + consumes ``charge_spin`` either way, but the freeze reads the MODEL's + ``get_dim_chg_spin()``. While the composition failed to forward it, + ``dim_chg_spin`` was 0 in metadata -- so the exported ABI had no + charge_spin slot and the C++ feeder never supplied one, making the + artifact disagree with its own eager model. Metadata-only, so no + inductor compile is needed. + """ + from deepmd.pt_expt.utils.serialization import ( + _collect_metadata, + ) + + config = copy.deepcopy(ZBL_CONFIG) + config["descriptor"]["add_chg_spin_ebd"] = True + model = get_model(config).to(torch.device("cpu")).eval() + meta = _collect_metadata(model, is_spin=False, lower_kind="graph") + assert meta["dim_chg_spin"] > 0, ( + "the bridged model's metadata dropped charge_spin; the exported " + "artifact would silently ignore the FiLM conditioning" + ) + + plain = copy.deepcopy(ZBL_CONFIG) + plain["descriptor"]["add_chg_spin_ebd"] = True + for key in ("bridging_method", "bridging_r_inner", "bridging_r_outer"): + plain.pop(key, None) + plain_model = get_model(plain).to(torch.device("cpu")).eval() + plain_meta = _collect_metadata(plain_model, is_spin=False, lower_kind="graph") + assert meta["dim_chg_spin"] == plain_meta["dim_chg_spin"] + + +def test_pair_exclusion_suppresses_the_analytical_term() -> None: + """Exclusion must remove the analytical term too, not just the learned one. + + Model-level ``pair_exclude_types`` is a neighbor-graph BUILD transform, + and both children of the bridged composition read the same graph. So an + excluded pair type must contribute neither a learned nor a ZBL term. + + This is not a free-standing preference: the composition is what drives + the build (and what the freeze metadata reads). Built without the + exclusion forwarded, the graph kept the excluded pairs and ZBL went on + interacting through them -- a large, silent error, since ZBL dominates at + short range. + + Fixture: an Ni-O dimer at 0.9 A, where (0, 1) is the ONLY pair present, + so excluding it must remove the analytical term entirely. + """ + cpu = torch.device("cpu") + coord = torch.tensor( + [[0.0, 0.0, 0.0], [0.9, 0.0, 0.0]], dtype=torch.float64, device=cpu + ).reshape(1, -1) + atype = torch.tensor([[0, 1]], dtype=torch.int64, device=cpu) + box = (torch.eye(3, dtype=torch.float64, device=cpu) * 20.0).reshape(1, 9) + + def energy(*, bridged: bool, excluded: bool) -> float: + config = copy.deepcopy(ZBL_CONFIG) + if not bridged: + for key in ("bridging_method", "bridging_r_inner", "bridging_r_outer"): + config.pop(key, None) + if excluded: + config["pair_exclude_types"] = [[0, 1]] + model = get_model(config).to(torch.device("cpu")).eval() + out = model(coord, atype, box=box)["energy"] + return float(out.detach().cpu().numpy().reshape(-1)[0]) + + # anti-vacuity: without exclusion the analytical term must dominate, + # otherwise the suppression assertion below would hold trivially. + zbl_contribution = energy(bridged=True, excluded=False) - energy( + bridged=False, excluded=False + ) + assert zbl_contribution > 1.0, ( + f"ZBL contributes only {zbl_contribution} at 0.9 A; the fixture no " + "longer exercises the analytical term" + ) + + # ... and with the pair type excluded, the bridged model must fall back + # EXACTLY onto the unbridged one: no ZBL, not merely less ZBL. + assert energy(bridged=True, excluded=True) == energy( + bridged=False, excluded=True + ), "the analytical ZBL term survived a pair-type exclusion" diff --git a/source/tests/pt_expt/test_dp_freeze.py b/source/tests/pt_expt/test_dp_freeze.py index cdaa4a02c6..b10e22d51a 100644 --- a/source/tests/pt_expt/test_dp_freeze.py +++ b/source/tests/pt_expt/test_dp_freeze.py @@ -134,43 +134,62 @@ def test_freeze_default_suffix(self) -> None: self.assertTrue(os.path.exists(expected)) def test_freeze_output_suffix_by_lower_kind(self) -> None: - """main() defaults a suffix-less output to .pt2 for --lower-kind graph - and .pte for nlist, while preserving an explicit .pte/.pt2 (iProzd - review). freeze() is mocked so the suffix logic is checked without the - AOTInductor compile cost. + """A suffix-less output defaults to .pt2 for lower_kind='graph' and + .pte for nlist, while preserving an explicit .pte/.pt2 (iProzd + review). The suffix follows the RESOLVED lower kind inside freeze() + (native-spin models force 'graph', so the CLI cannot pick it before + the model is built); the mapping is checked on the helper freeze() + defers to. End-to-end application through main()/freeze() is covered + by test_freeze_default_suffix (nlist) and + test_native_spin_default_freeze_routes_to_graph in test_dpa4_export + (graph). """ - from unittest import ( - mock, + from deepmd.pt_expt.entrypoints.main import ( + _default_output_path, ) cases = [ - ("graph", "out_g", None, ".pt2"), # graph, no suffix -> .pt2 - ("nlist", "out_n", None, ".pte"), # nlist, no suffix -> .pte - ("graph", "out_g_explicit", ".pte", ".pte"), # explicit .pte kept - ("nlist", "out_n_explicit", ".pt2", ".pt2"), # explicit .pt2 kept + ("graph", "out_g", ".pt2"), # graph, no suffix -> .pt2 + ("nlist", "out_n", ".pte"), # nlist, no suffix -> .pte + ("graph", "out_g_explicit.pte", ".pte"), # explicit .pte kept + ("nlist", "out_n_explicit.pt2", ".pt2"), # explicit .pt2 kept ] - for lower_kind, stem, explicit, expected_suffix in cases: - with self.subTest(lower_kind=lower_kind, explicit=explicit): - name = stem + (explicit or "") - captured: dict = {} - - def _fake_freeze(model, output, head=None, lower_kind="nlist", **kw): - captured["output"] = output - captured["lower_kind"] = lower_kind - - flags = argparse.Namespace( - command="freeze", - checkpoint_folder=self.ckpt_file, - output=os.path.join(self.tmpdir, name), - head=None, - lower_kind=lower_kind, - log_level=2, - log_path=None, + for lower_kind, name, expected_suffix in cases: + with self.subTest(lower_kind=lower_kind, name=name): + resolved = _default_output_path( + os.path.join(self.tmpdir, name), lower_kind ) - with mock.patch("deepmd.pt_expt.entrypoints.main.freeze", _fake_freeze): - main(flags) - self.assertTrue(captured["output"].endswith(expected_suffix)) - self.assertEqual(captured["lower_kind"], lower_kind) + self.assertTrue(resolved.endswith(expected_suffix)) + + def test_freeze_main_passes_lower_kind_through(self) -> None: + """main() forwards --lower-kind and the raw output path to freeze() + (suffix defaulting is owned by freeze(), after native-spin + resolution). + """ + from unittest import ( + mock, + ) + + captured: dict = {} + + def _fake_freeze(model, output, head=None, lower_kind="nlist", **kw): + captured["output"] = output + captured["lower_kind"] = lower_kind + + raw_output = os.path.join(self.tmpdir, "out_passthrough") + flags = argparse.Namespace( + command="freeze", + checkpoint_folder=self.ckpt_file, + output=raw_output, + head=None, + lower_kind="graph", + log_level=2, + log_path=None, + ) + with mock.patch("deepmd.pt_expt.entrypoints.main.freeze", _fake_freeze): + main(flags) + self.assertEqual(captured["output"], raw_output) + self.assertEqual(captured["lower_kind"], "graph") def test_freeze_graph_rejects_ineligible(self) -> None: """``--lower-kind graph`` on a non-graph-eligible model (se_e2_a, diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py index 9e83f8e135..193fa81911 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -762,6 +762,9 @@ def has_default_fparam(self) -> bool: def get_default_fparam(self) -> list[float]: return [0.0, 1.0] + def has_spin(self) -> bool: + return False + def has_chg_spin_ebd(self) -> bool: return False @@ -796,6 +799,9 @@ def has_default_fparam(self) -> bool: def get_default_fparam(self) -> None: return None + def has_spin(self) -> bool: + return False + def has_chg_spin_ebd(self) -> bool: return False diff --git a/source/tests/pt_expt/test_wrapper.py b/source/tests/pt_expt/test_wrapper.py index 594067cd39..da2b500cc9 100644 --- a/source/tests/pt_expt/test_wrapper.py +++ b/source/tests/pt_expt/test_wrapper.py @@ -22,6 +22,10 @@ def __init__(self, *, fail_forward: bool = False) -> None: self.fail_forward = fail_forward self.last_requires_grad: tuple[bool, ...] | None = None + def has_spin(self) -> bool: + """Mirror the base-model capability contract (concrete, no getattr probe).""" + return False + def forward( self, coord: torch.Tensor, diff --git a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py index 8dd8af74c1..bcd0b3b0c8 100644 --- a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py +++ b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py @@ -19,6 +19,7 @@ from deepmd.pt_expt.utils.serialization import ( _graph_edge_dtype, + _needs_with_comm_artifact, _supports_graph_export, deserialize_to_file, ) @@ -352,3 +353,68 @@ def test_graph_trace_version_guard_dpa2_compact_pairs( }, "fitting_net": {"neuron": [8, 8], "seed": 1}, } + + +def _build_model(model_kind: str) -> torch.nn.Module: + """Build a small pt_expt model for ``_needs_with_comm_artifact`` tests. + + No AOTI compile is involved — the caller only inspects the returned + model's descriptor capability methods. + + Parameters + ---------- + model_kind : str + ``"dpa4"`` (bridging-free SeZM, config shared with + ``test_dpa4_export.py``) or ``"dpa2"`` (``DPA2_GUARD_CONFIG`` above). + + Returns + ------- + torch.nn.Module + The constructed pt_expt model, on CPU, in eval mode. + """ + from deepmd.pt_expt.model.get_model import get_model as get_pt_expt_model + + if model_kind == "dpa4": + from ..model.test_dpa4_export import ( + _DPA4_CONFIG, + ) + + config = _DPA4_CONFIG + elif model_kind == "dpa2": + config = DPA2_GUARD_CONFIG + else: + raise ValueError(f"unknown model_kind {model_kind!r}") + model = get_pt_expt_model(copy.deepcopy(config)) + model.to("cpu") + model.eval() + return model + + +@pytest.mark.parametrize( + "model_kind,lower_kind,expected", + [ + ("dpa4", "graph", True), # graph lower has real border exchange now + ( + "dpa4", + "nlist", + False, + ), # dense lower is comm-less: no artifact, no trace crash + ( + "dpa2", + "nlist", + True, + ), # dense with-comm is dpa2's production MP path — unchanged + ("dpa2", "graph", True), # graph with-comm unchanged + ], +) +def test_needs_with_comm_artifact_kind_aware(model_kind, lower_kind, expected) -> None: + """``_needs_with_comm_artifact`` is lower-kind-aware for DPA4, unchanged for dpa2. + + DPA4's graph lower carries a real per-layer ``border_op`` exchange, but + its dense (nlist) lower adapter raises on ``comm_dict`` — so the dense + kind must not request a with-comm artifact (it would crash the trace). + dpa2 implements comm on both lowers (no ``dense_lower_supports_comm`` + override), so both kinds stay ``True``. + """ + model = _build_model(model_kind) + assert _needs_with_comm_artifact(model, lower_kind) is expected diff --git a/source/tests/tf2/test_training.py b/source/tests/tf2/test_training.py index dc6324684c..e282e983db 100644 --- a/source/tests/tf2/test_training.py +++ b/source/tests/tf2/test_training.py @@ -438,8 +438,11 @@ def _input_type_cast( fparam: Any = None, aparam: Any = None, charge_spin: Any = None, - ) -> tuple[Any, Any, Any, Any, Any, Any]: - return coord, box, fparam, aparam, charge_spin, coord.dtype + spin: Any = None, + ) -> tuple[Any, Any, Any, Any, Any, Any, Any]: + # 7-tuple: matches make_model's ``_input_type_cast`` (the ``spin`` + # slot for the native-spin graph route; tf2 discards it). + return coord, box, fparam, aparam, charge_spin, spin, coord.dtype def _output_type_cast( self, diff --git a/source/tests/universal/common/cases/descriptor/utils.py b/source/tests/universal/common/cases/descriptor/utils.py index 96eee7b757..03dd5248da 100644 --- a/source/tests/universal/common/cases/descriptor/utils.py +++ b/source/tests/universal/common/cases/descriptor/utils.py @@ -27,6 +27,28 @@ class DescriptorTestCase(TestCaseSingleFrameWithNlist): def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) + def test_capability_contract(self) -> None: + """Pin the ``BaseDescriptor`` capability-method contract. + + Every capability query is declared on ``make_base_descriptor`` with a + concrete default (never probed via ``getattr(..., )`` duck + typing), so ALL descriptors must answer direct method calls. + Eligibility values differ per descriptor/backend, so the assertions + are invariant-based; descriptor-specific ``True`` branches are pinned + in the implementing descriptors' own test files. + """ + assert isinstance(self.module.uses_graph_lower(), bool) + assert isinstance(self.module.uses_compact_edge_pairs(), bool) + assert isinstance(self.module.supports_native_spin(), bool) + assert isinstance(self.module.supports_charge_spin(), bool) + # Table hook must be callable on every descriptor (``None`` when the + # descriptor embeds types internally or has no graph lower). + self.module.graph_type_embedding_table() + # Pinned postcondition: after the escape hatch, the graph lower is + # off -- a no-op on descriptors without one. + self.module.disable_graph_lower() + assert self.module.uses_graph_lower() is False + def test_forward_consistency(self) -> None: ret = [] for module in self.modules_to_test: diff --git a/source/tests/universal/common/cases/model/utils.py b/source/tests/universal/common/cases/model/utils.py index 7728aea895..5ff4254934 100644 --- a/source/tests/universal/common/cases/model/utils.py +++ b/source/tests/universal/common/cases/model/utils.py @@ -117,6 +117,20 @@ def test_has_message_passing(self) -> None: module.has_message_passing(), self.expected_has_message_passing ) + def test_has_spin(self) -> None: + """Test has_spin. + + Declared on the base model with a concrete default of ``False`` + (direct method call, never a ``getattr`` probe); spin model wrappers + override it to ``True`` -- which is exactly the suite's ``test_spin`` + flag. Checked on the raw module only (not the jit-scripted copies in + ``modules_to_test``): the capability contract is a Python-API + contract, and pt's TorchScript models do not ``@torch.jit.export`` + it. + """ + expected = getattr(self, "test_spin", False) + self.assertEqual(self.module.has_spin(), expected) + def test_forward(self) -> None: """Test forward and forward_lower.""" test_spin = getattr(self, "test_spin", False)