diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 3a7922e8bd..25b770a5fb 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -199,6 +199,28 @@ def uses_graph_lower(self) -> bool: """ return False + def fused_decomposition(self) -> "tuple[Any, Any] | None": + """Returns the parts a fused energy-force pipeline evaluates, if any. + + Such a pipeline computes one descriptor-fitting atomic model and, + inside the same edge scan, one analytical pair potential. A model that + is exactly their plain sum answers with ``(learned, pair_potential)``, + either entry ``None`` when absent; every other model answers ``None`` + (the concrete default) and keeps the autograd lower. + + Generic capability in the sense of :meth:`uses_graph_lower`: the model + layer asks instead of reaching into the atomic model for a descriptor, + so the answer stays correct for compositions. + + Returns + ------- + tuple[Any, Any] or None + The learned descriptor-fitting atomic model and the analytical + pair potential, either ``None`` when absent, or ``None`` for a + model that is not such a sum. + """ + return None + def supports_native_spin(self) -> bool: """Returns whether this atomic model consumes a per-atom spin input. diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index ffd33b5bd2..b9f20ed8c7 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -213,6 +213,10 @@ def uses_graph_lower(self) -> bool: """Delegates to this model's own descriptor.""" return bool(self.descriptor.uses_graph_lower()) + def fused_decomposition(self) -> "tuple[DPAtomicModel, None]": + """A descriptor-fitting model is its own learned part.""" + return self, None + def has_message_passing_across_ranks(self) -> bool: """Delegates to this model's own descriptor.""" return bool(self.descriptor.has_message_passing_across_ranks()) @@ -369,7 +373,9 @@ def forward_atomic( ) charge_spin = xp.tile(xp.reshape(cs_array, (1, -1)), (nframes, 1)) - descriptor, rot_mat, g2, h2, sw = self.descriptor( + # A bridged descriptor returns a sixth element, the per-atom source + # gate of its bridging window; every other descriptor returns five. + outputs = self.descriptor( extended_coord, extended_atype, nlist, @@ -377,11 +383,16 @@ def forward_atomic( comm_dict=comm_dict, charge_spin=charge_spin if self.add_chg_spin_ebd else None, ) + descriptor, rot_mat, g2, h2 = outputs[:4] + node_gate = outputs[5] if len(outputs) > 5 else None # The vacuum descriptor of every type is handed to a fitting that - # references its atoms; other fittings do not take the keyword. + # references its atoms; other fittings do not take the keyword. The + # source gate fades the learned part of the fitting output: an atom + # of a frozen pair keeps only the bias of its type. vacuum_kwargs = {} if self.fitting_net.needs_vacuum_descriptor(): vacuum_kwargs["vacuum_descriptor"] = self.vacuum_descriptor() + gate_kwargs = {} if node_gate is None else {"node_gate": node_gate} ret = self.fitting_net( descriptor, atype, @@ -391,6 +402,7 @@ def forward_atomic( fparam=fparam, aparam=aparam, **vacuum_kwargs, + **gate_kwargs, ) return ret @@ -401,7 +413,11 @@ def vacuum_descriptor(self) -> Array: conditioned as the neutral ground-state atom: zero charge with the ground-state multiplicity when the descriptor takes the charge/spin condition, and a spin vector of one Bohr magneton per unpaired - electron when it takes the native spin. + electron when it takes the native spin. A descriptor with a graph + lower is evaluated on the neighbor graph of those single-atom frames, + which holds no real edge; a dense descriptor on a neighbor list of its + own capacity filled with padding, since its equations are sized by + that capacity. Returns ------- @@ -417,6 +433,35 @@ def vacuum_descriptor(self) -> Array: atype = xp.reshape( xp.arange(ntypes, dtype=xp.int64, device=device), (ntypes, 1) ) + conditions = { + name: xp.asarray(table, dtype=bias.dtype, device=device) + for name, table in self.vacuum_conditions().items() + } + if self.descriptor.uses_graph_lower(): + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, + ) + + graph = build_neighbor_graph( + coord, atype, None, self.descriptor.get_rcut(), with_csr=True + ) + # Only descriptors whose ``call_graph`` declares a condition take + # the keyword; a declared condition the model does not use is None. + spin_kwargs = ( + {"spin": conditions.get("spin")} if self.supports_native_spin() else {} + ) + charge_spin_kwargs = ( + {"charge_spin": conditions.get("charge_spin")} + if self.supports_charge_spin + else {} + ) + return self.descriptor.call_graph( + graph, + xp.reshape(atype, (ntypes,)), + type_embedding=self.descriptor.graph_type_embedding_table(), + **spin_kwargs, + **charge_spin_kwargs, + )[0] nlist = xp.full( (ntypes, 1, self.descriptor.get_nnei()), -1, @@ -424,10 +469,6 @@ def vacuum_descriptor(self) -> Array: device=device, ) mapping = xp.zeros((ntypes, 1), dtype=xp.int64, device=device) - conditions = { - name: xp.asarray(table, dtype=bias.dtype, device=device) - for name, table in self.vacuum_conditions().items() - } if "spin" in conditions: conditions["spin"] = conditions["spin"][:, None, :] descriptor = self.descriptor( @@ -617,7 +658,9 @@ def forward_atomic_graph( charge_spin_kwargs = ( {"charge_spin": charge_spin} if self.supports_charge_spin else {} ) - gg, rot_mat = self.descriptor.call_graph( + # A bridged descriptor returns a third element, the per-node source + # gate of its bridging window; every other descriptor returns two. + outputs = self.descriptor.call_graph( graph, atype, type_embedding=type_embedding, @@ -625,13 +668,21 @@ def forward_atomic_graph( **spin_kwargs, **charge_spin_kwargs, ) + gg, rot_mat = outputs[0], outputs[1] + node_gate = outputs[2] if len(outputs) > 2 else None vacuum_kwargs = {} if vacuum_ref: vacuum_kwargs["vacuum_descriptor"] = gg[n_real:] gg = gg[:n_real] atype = atype[:n_real] rot_mat = None if rot_mat is None else rot_mat[:n_real] + node_gate = None if node_gate is None else node_gate[:n_real] # === Step 4. Fitting === + # The source gate of a bridged descriptor fades the learned part of + # the fitting output: an atom of a frozen pair keeps only the bias of + # its type, so along the closing of a pair the learned energy follows + # the switch linearly while the analytical pair term takes over. + gate_kwargs = {} if node_gate is None else {"node_gate": node_gate} return self.fitting_net.call_graph( gg, atype, @@ -641,6 +692,7 @@ def forward_atomic_graph( fparam=fparam_node, aparam=aparam_node, **vacuum_kwargs, + **gate_kwargs, ) def compute_or_load_stat( diff --git a/deepmd/dpmodel/atomic_model/inner_potential.py b/deepmd/dpmodel/atomic_model/inner_potential.py index 4e2b727688..4d3c51f2d7 100644 --- a/deepmd/dpmodel/atomic_model/inner_potential.py +++ b/deepmd/dpmodel/atomic_model/inner_potential.py @@ -1,10 +1,15 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Analytical pair potentials for Zone bridging (backend-agnostic port of -``deepmd.pt``'s ``InnerPotential``). Lives in the atomic-model package: -the atomic layer owns per-atom energy assembly, where the ZBL term is -injected on the graph route. +"""Analytical pair potentials for Zone bridging, defined once for every backend. + +Lives in the atomic-model package: the atomic layer owns per-atom energy +assembly, where the analytical term is injected on the graph route. """ +import functools +import json +from pathlib import ( + Path, +) from typing import ( Any, ) @@ -62,17 +67,98 @@ _KE_EV_A = 14.3996 # Coulomb constant in eV·Å _A_BOHR = 0.5291772109 # Bohr radius in Å +# Coefficients of the NLH screening function, one row per unordered element +# pair, in the file beside this module. +# +# Provenance and attribution. The reference data is the open-data package +# K. Nordlund, G. Hobler and S. Lehtola, "Data sets for the publication +# 'Repulsive interatomic potentials calculated at three levels of theory'", +# Zenodo version 1.0, 16 November 2024, doi 10.5281/zenodo.14172633, +# distributed under the Creative Commons Attribution 4.0 International licence +# (CC BY 4.0, https://creativecommons.org/licenses/by/4.0/). The functional +# form is that of K. Nordlund, S. Lehtola and G. Hobler, Phys. Rev. A 111, +# 032818 (2025), doi 10.1103/PhysRevA.111.032818, with its erratum +# Phys. Rev. A 112, 059901 (2025), doi 10.1103/cdrk-x7my, published under the +# same licence. +# +# The shipped numbers are MODIFIED material: they are this project's own refit +# of that open data and are not the coefficients published by Nordlund, +# Lehtola and Hobler. The fit imposes conditions the published one does not +# carry -- non-negative amplitudes, amplitudes summing to one, and a bound on +# the pair energy at long range -- because the term is added on every neighbour +# pair inside the cutoff without a switching function. Two groups of rows come +# from a different reference than the rest: B-Ne (Z1 = 5, Z2 = 10) is fitted to +# the package's Hartree-Fock MP2 screening function, and every pair containing +# an element with Z > 92, which the reference data does not cover, is fitted to +# the Ziegler-Biersack-Littmark universal potential. Neither the authors of the +# reference data nor the American Physical Society endorse this project. +# +# The same record, machine-readable, is in the file's ``meta`` entry and is +# returned by :func:`nlh_provenance`. +NLH_COEFFICIENTS_FILE = Path(__file__).with_name("nlh_coefficients.npz") + +_SERIES_TERMS = 4 + + +@functools.lru_cache(maxsize=1) +def _nlh_lookup() -> tuple[Array, Array]: + r"""Load the NLH screening coefficients into charge-indexed lookups. + + Returns + ------- + amplitude : Array + Amplitudes :math:`a_k` with shape ``(Z_max + 1, Z_max + 1, 4)``, + unitless and summing to one along the last axis; symmetric in the two + charge axes, with row and column zero unused. + rate : Array + Decay rates :math:`b_k` with the same shape, in Å⁻¹. + """ + with np.load(NLH_COEFFICIENTS_FILE, allow_pickle=False) as data: + pairs, a, b = data["pairs"], data["a"], data["b"] + # The file stores every unordered pair once, so the largest charge it + # names is the highest index either axis has to address. + z_max = int(pairs.max()) + amplitude = np.zeros((z_max + 1, z_max + 1, _SERIES_TERMS), dtype=np.float64) + rate = np.zeros_like(amplitude) + z1, z2 = pairs[:, 0].astype(np.int64), pairs[:, 1].astype(np.int64) + amplitude[z1, z2] = amplitude[z2, z1] = a + rate[z1, z2] = rate[z2, z1] = b + return amplitude, rate + + +def nlh_provenance() -> dict[str, Any]: + """Read the provenance recorded inside the shipped coefficient file. + + Returns + ------- + dict + The reference data and article the coefficients derive from, the + conditions and objective of the fit, and which reference covers which + rows. + """ + with np.load(NLH_COEFFICIENTS_FILE, allow_pickle=False) as data: + return json.loads(str(data["meta"])) + class InnerPotential(NativeOP): - """Analytical pair potential for Zone bridging. + r"""Analytical pair potential for Zone bridging. + + Every supported potential is a screened Coulomb series in the pair + separation, + + .. math:: + + V_{ab}(r)=\frac1r\sum_{k=1}^{4}A_{abk}\,e^{-c_{abk}r}, - 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.InnerPotential``. + so one table of eight constants per ordered type pair describes all of + them and both evaluation routes read it: :attr:`series_table` carries it in + double precision for :meth:`call`, and :attr:`pair_table` is its float32 + view, the layout the fused kernels require. The term is 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(r_ij) / 2`` to both + atom i and atom j, avoiding double-counting from the symmetric neighbor + list. The class is written against the array API, so every backend + evaluates the same definition on its own array type. Parameters ---------- @@ -80,7 +166,9 @@ class InnerPotential(NativeOP): 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. + Potential formula: ``"zbl"`` for the Ziegler-Biersack-Littmark + universal potential, ``"nlh"`` for the bundled Nordlund-Lehtola-Hobler + coefficients. Case-insensitive. Raises ------ @@ -89,32 +177,135 @@ class InnerPotential(NativeOP): not found in the periodic table. """ + # Rebuilt from ``type_map`` and ``mode`` by ``__init__`` and + # ``change_type_map``, so a wrapped backend keeps them out of its + # checkpoints. + CONFIG_DERIVED_ARRAYS = ("pair_table", "series_table") + + SUPPORTED_MODES = ("ZBL", "NLH") + def __init__(self, type_map: list[str], mode: str = "zbl") -> None: super().__init__() mode = str(mode).upper() - if mode != "ZBL": + if mode not in self.SUPPORTED_MODES: raise ValueError(f"Unknown InnerPotential 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) + self.ntypes = len(type_map) + table = self.series_table_from_type_map(type_map, mode=mode) + self.series_table = table + self.pair_table = table.astype(np.float32) - @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. + @classmethod + def series_table_from_type_map( + cls, + type_map: list[str], + mode: str = "ZBL", + dtype: Any = np.float64, + like: Array | None = None, + ) -> Array: + r"""Tabulate the screened Coulomb series of every ordered type pair. + + Both modes share the series + + .. math:: + + V_{ab}(r)=\frac1r\sum_{k=1}^{4}A_{abk}\,e^{-c_{abk}r},\qquad + A_{abk}=k_e Z_aZ_b\,a_{abk}, + + and differ only in where the unitless amplitudes :math:`a_{abk}` and + the decay rates :math:`c_{abk}` come from. ``"ZBL"`` takes the + universal screening function, whose amplitudes are the same for every + pair and whose rates are the universal exponents divided by the + pair's screening length :math:`a_{ab}=0.88534\,a_0/(Z_a^{0.23}+ + Z_b^{0.23})`. ``"NLH"`` reads both from the bundled per-pair + coefficients, where the rates are already in Å⁻¹. + + A kernel therefore evaluates the energy and its radial derivative from + eight per-pair constants and four exponentials in either mode. Parameters ---------- type_map : list[str] Element symbols; index corresponds to ``atype`` values. + mode : str + Potential formula, case-insensitive. + dtype + Floating-point dtype of the result. The fused kernels read the + table in single precision, :meth:`call` in double. 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 + When given, the result is created in this array's namespace 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 + Coefficients ``[A_1..A_4, c_1..c_4]`` with shape + ``((T + 1) ** 2, 8)``, in eV Å and Å⁻¹. Row + ``center * (T + 1) + neighbor`` follows the ordered pair index of + the fused descriptor caches; rows of the padding type ``T`` vanish. + """ + mode = str(mode).upper() + if mode not in cls.SUPPORTED_MODES: + raise ValueError(f"Unknown InnerPotential mode: {mode}") + charge = cls._lookup_from_type_map(type_map) + ntypes = charge.shape[0] + if mode == "NLH": + lut_a, lut_b = _nlh_lookup() + index = charge.astype(np.int64) + amplitude = lut_a[index[:, None], index[None, :]] + rate = lut_b[index[:, None], index[None, :]] + else: + screening = ( + 0.88534 * _A_BOHR / (charge[:, None] ** 0.23 + charge[None, :] ** 0.23) + ) + amplitude = np.broadcast_to( + np.asarray(_ZBL_A_COEFF), (ntypes, ntypes, _SERIES_TERMS) + ) + rate = np.asarray(_ZBL_B_COEFF) / screening[:, :, None] + table = np.zeros((ntypes + 1, ntypes + 1, 2 * _SERIES_TERMS), dtype=np.float64) + table[:ntypes, :ntypes, :_SERIES_TERMS] = ( + _KE_EV_A * (charge[:, None] * charge[None, :])[:, :, None] * amplitude + ) + table[:ntypes, :ntypes, _SERIES_TERMS:] = rate + table = table.reshape(-1, 2 * _SERIES_TERMS).astype(dtype, copy=False) + return cls._in_namespace_of(table, like) + + @staticmethod + def _in_namespace_of(table: np.ndarray, like: Array | None) -> Array: + """Place a NumPy table in the namespace and device of another array. + + Parameters + ---------- + table : np.ndarray + The table to place. + like : Array, optional + The array whose namespace and device the result adopts. ``None`` + returns the table unchanged. + + Returns + ------- + Array + The table, in ``like``'s namespace and on its device. + """ + if like is None: + return table + xp = array_api_compat.array_namespace(like) + return xp.asarray(table, device=array_api_compat.device(like)) + + @staticmethod + def _lookup_from_type_map(type_map: list[str]) -> np.ndarray: + """Build the per-type nuclear-charge lookup from element symbols. + + Parameters + ---------- + type_map : list[str] + Element symbols; index corresponds to ``atype`` values. + + Returns + ------- + np.ndarray Nuclear charges, shape ``(len(type_map),)``. Raises @@ -128,21 +319,17 @@ def _lookup_from_type_map(type_map: list[str], like: Array | None = None) -> Arr 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)) + return np.asarray(atomic_numbers, dtype=np.float64) 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 + """Rebuild the coefficient tables for a new type map. + + This class owns the coefficient tables, so it owns every update of + them: the symbols, their count and the two tables are one piece of + state and are replaced together. Reordering, adding and dropping + elements are all covered, because a table is rebuilt from the symbols + rather than permuted, so no index bookkeeping can drift. Each rebuilt + array keeps the current one's namespace and device, so a wrapped backend (pt_expt buffer on CPU or CUDA) is updated in place. Parameters @@ -150,39 +337,39 @@ def change_type_map(self, type_map: list[str]) -> None: type_map : list[str] The new element symbols. """ - self.atomic_numbers = self._lookup_from_type_map( - type_map, like=self.atomic_numbers + table = self.series_table_from_type_map(type_map, mode=self.mode) + self.series_table = self._in_namespace_of(table, self.series_table) + self.pair_table = self._in_namespace_of( + table.astype(np.float32), self.pair_table ) self.type_map = list(type_map) - self.ntypes_real = len(type_map) + self.ntypes = 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. + def _series_energy(xp: Any, r: Array, row: Array) -> Array: + r"""Evaluate the screened Coulomb series from its per-pair constants. + + .. math:: + + V(r)=\frac1r\sum_{k=1}^{4}A_k\,e^{-c_k r} Parameters ---------- xp - The array namespace of ``r``/``zi``/``zj``. + The array namespace of ``r`` and ``row``. 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 (...). + Pair distances with shape (E,) in Å. + row : Array + Series constants ``[A_1..A_4, c_1..c_4]`` with shape (E, 8), in + eV Å and Å⁻¹. Returns ------- Array - Pair energies with shape (...) in eV. + Pair energies with shape (E,) 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 + decay = xp.exp(-row[:, _SERIES_TERMS:] * r[:, None]) + return xp.sum(row[:, :_SERIES_TERMS] * decay, axis=-1) / r def call( self, @@ -191,16 +378,16 @@ def call( 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. + """Scatter per-edge analytical 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). + yields the analytical force/virial through the shared edge + backward). edge_index : Array (2, E) ``[src, dst]`` edge endpoints (flat node indices). atype_flat : Array @@ -209,53 +396,32 @@ def call( (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 + Per-atom analytical 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 + stride = self.ntypes + 1 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, + # Padded atoms carry type -1; their edges are masked below, so the + # table lookup only needs a valid index for them. + atype_row = xp.clip(xp.astype(atype_flat, xp.int64), min=0) + # Row ``center * (T + 1) + neighbor``, the ordered-pair index the + # fused kernels use, with the destination as the center. + row = xp.take( + xp_asarray_nodetach(xp, self.series_table, dtype=xp.float64, device=device), + xp.take(atype_row, dst, axis=0) * stride + xp.take(atype_row, src, axis=0), 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) + pair_e = self._series_energy(xp, r, row) + pair_e = pair_e * xp.astype(xp.astype(edge_mask, xp.bool), 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 ( @@ -285,7 +451,7 @@ class InnerPotentialAtomicModel(BaseAtomicModel): type_map : list[str] Element symbols; index corresponds to ``atype`` values. mode : str - Potential formula (currently ``"zbl"``). + Potential formula, ``"zbl"`` or ``"nlh"``. Case-insensitive. rcut : float Cut-off radius this model declares (the composition uses the max over children; pass the learned model's). @@ -396,6 +562,30 @@ def uses_graph_lower(self) -> bool: """ return True + def graph_edge_dtype(self) -> str: + """The term accepts single-precision edges, so under the composition + rule the learned sibling decides the dtype of the shared edge tensor. + """ + return "float32" + + def fused_decomposition(self) -> tuple[None, InnerPotential]: + """The term is a pair potential with no learned part.""" + return None, self.potential + + def enable_compression( + self, + min_nbor_dist: float, + table_extrapolate: float = 5, + table_stride_1: float = 0.01, + table_stride_2: float = 0.1, + check_frequency: int = -1, + ) -> None: + """Analytical term: nothing to tabulate.""" + + def compression_needs_min_nbor_dist(self) -> bool: + """Analytical term: compression reads no neighbor statistics.""" + return False + def forward_atomic( self, *args: Any, @@ -446,7 +636,6 @@ def forward_atomic_graph( atype, graph.edge_mask, n_node=n_node, - real_type_count=len(self.type_map), ) return {"energy": xp.reshape(energy, (n_node, 1))} diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 5d879043af..e224cbdc2b 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -336,6 +336,48 @@ def uses_graph_lower(self) -> bool: """ return all(m.uses_graph_lower() for m in self.models) + def fused_decomposition(self) -> "tuple[Any, Any] | None": + """Merge the children's parts when the composition is their plain sum. + + A fused pipeline evaluates one learned model plus at most one pair + potential on the parent's node axis, under the parent's masks and + output bias. That describes this composition only with ``"sum"`` + weights, shared type maps, and children that carry no exclusion of + their own; a bridged ``[learned, inner_potential]`` pair qualifies. + The analytical child contributes no output bias to the pipeline, + because its output statistics are never computed and remain zero. + """ + if self.weights != "sum" or not self._graph_mapping_is_identity: + return None + parts = [model.fused_decomposition() for model in self.models] + if any(part is None for part in parts) or any( + model.atom_excl is not None or model.pair_excl is not None + for model in self.models + ): + return None + learned = [part[0] for part in parts if part[0] is not None] + potentials = [part[1] for part in parts if part[1] is not None] + if len(learned) != 1 or len(potentials) > 1: + return None + return learned[0], potentials[0] if potentials else None + + def compute_fitting_input_stats( + self, sampled_func: Callable[[], list[dict]] + ) -> None: + """Recompute the input statistics of every learned sub-model's fitting net. + + Analytical and tabulated sub-models carry no fitting net and are left + alone. + + Parameters + ---------- + sampled_func + The lazy sampled function to get data frames from different data systems. + """ + for model in self.models: + if isinstance(model, DPAtomicModel): + model.fitting_net.compute_input_stats(sampled_func) + def supports_native_spin(self) -> bool: """Spin-capable when ANY child consumes the spin input. diff --git a/deepmd/dpmodel/atomic_model/nlh_coefficients.npz b/deepmd/dpmodel/atomic_model/nlh_coefficients.npz new file mode 100644 index 0000000000..e5d5a32bbc Binary files /dev/null and b/deepmd/dpmodel/atomic_model/nlh_coefficients.npz differ diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index 59c027bbf9..7e264fb3e9 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -74,9 +74,18 @@ from deepmd.dpmodel.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.bridging import ( + BRIDGING_RECORD_VERSION, + check_bridging_record_version, + check_window_inside_cutoff, + migrate_inner_clamp_keys, +) from deepmd.utils.charge_state import ( validate_charge_state, ) +from deepmd.utils.element_radii import ( + contact_radius_table, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -117,6 +126,7 @@ InnerClamp, RadialBasis, RadialMLP, + pair_contact, ) from .dpa4_nn.utils import ( ATTN_RES_MODES, @@ -582,12 +592,21 @@ class DescrptDPA4(NativeOP, BaseDescriptor): Random seed(s). type_map Type names. - inner_clamp_r_inner - Inner radius for distance saturation in Å. If both inner and outer radii - are set, the descriptor freezes short-range descriptor geometry inside - the zone-bridging window. - inner_clamp_r_outer - Outer radius for distance saturation in Å. + inner_clamp_f_inner + Inner radius of the zone-bridging window, as a fraction of each pair's + own length scale. When both fractions are set, the descriptor freezes + short-range geometry inside the window. Given together with + ``inner_clamp_f_outer`` or not at all. + inner_clamp_f_outer + Outer radius of the zone-bridging window, as a fraction of each pair's + own length scale. At or beyond it an edge contributes exactly as it + does without a window. + inner_clamp_scale + Length scale the two fractions measure against. ``"covalent"`` sizes + every pair by the sum of the covalent radii of its two elements, so one + window shape serves every element combination and a ``type_map`` is + required; ``"absolute"`` gives every element the unit radius, which + makes the pair scale 1 Å and the two fractions radii in Å. add_chg_spin_ebd If True, add frame-level charge/spin condition embedding to scalar type features before edge features are built. @@ -605,7 +624,7 @@ class DescrptDPA4(NativeOP, BaseDescriptor): """ _ENV_DIM: int = 1 # Use se_r style (radial only) for EnvMatStatSe compatibility - LATEST_VERSION: float = 1.2 + LATEST_VERSION: float = 1.3 def __init__( self, @@ -668,12 +687,12 @@ def __init__( trainable: bool = True, seed: int | list[int] | None = None, type_map: list[str] | None = None, - inner_clamp_r_inner: float | None = None, - inner_clamp_r_outer: float | None = None, + inner_clamp_f_inner: float | None = None, + inner_clamp_f_outer: float | None = None, + inner_clamp_scale: str = "covalent", add_chg_spin_ebd: bool = False, default_chg_spin: list[float] | None = None, use_spin: list[bool] | None = None, - **kwargs: Any, ) -> None: self.version = float(self.LATEST_VERSION) self._graph_lower_disabled = False @@ -830,33 +849,58 @@ def __init__( # === Zone bridging: InnerClamp + Source Freeze Propagation Gate === # Both the geometry clamp (``InnerClamp``) and the message-passing # switch (``BridgingSwitch``) are activated together on the same - # ``[r_inner, r_outer]`` window. The clamp freezes scalar distance - # on every ``(j, k)`` edge with ``r_{jk} < r_inner``; the switch - # feeds a per-edge C3 amplitude into ``compute_edge_src_gate`` so - # that any node with a frozen neighbor cannot propagate - # information through the GNN, closing the direction / multi-hop - # leakage channels that a pure ``InnerClamp`` cannot reach. Both - # modules are parameter-free, so enabling bridging does not add - # any keys to the descriptor's state dict. - self.inner_clamp_r_inner = ( - float(inner_clamp_r_inner) if inner_clamp_r_inner is not None else None + # window, two fractions of the length scale each pair contributes + # through ``contact_radius``. The switch spans the whole window; the + # clamp freezes the distance the descriptor sees on every ``(j, k)`` + # edge two fifths into the window, just below the frame-filter point + # at the midpoint, so that the displayed distance still moves where + # the training data end and the network can fit the labels there. The + # switch feeds a per-edge C3 amplitude of the true length into + # ``compute_source_gates``, whose per-node product over every pair + # other than the edge's own is folded into the envelope of every edge + # the node emits, so that a node with a frozen neighbor is absent + # from every reduction of its other neighbors as if those edges were + # deleted, while the frozen pair keeps seeing itself at the clamped + # distance. The full per-node product also fades the learned atomic + # energy in the atomic model, so a frozen atom contributes the bias + # of its type alone. Both modules are parameter-free, so enabling + # bridging does not add any keys to the descriptor's state dict. + if (inner_clamp_f_inner is None) != (inner_clamp_f_outer is None): + raise ValueError( + "`inner_clamp_f_inner` and `inner_clamp_f_outer` must be given " + "together." + ) + if inner_clamp_scale not in ("covalent", "absolute"): + raise ValueError( + "`inner_clamp_scale` must be 'covalent' or 'absolute', got " + f"{inner_clamp_scale!r}." + ) + self.bridging_f_inner = ( + None if inner_clamp_f_inner is None else float(inner_clamp_f_inner) ) - self.inner_clamp_r_outer = ( - float(inner_clamp_r_outer) if inner_clamp_r_outer is not None else None + self.bridging_f_outer = ( + None if inner_clamp_f_outer is None else float(inner_clamp_f_outer) ) - if ( - self.inner_clamp_r_inner is not None - and self.inner_clamp_r_outer is not None - ): - self.inner_clamp: InnerClamp | None = InnerClamp( - self.inner_clamp_r_inner, self.inner_clamp_r_outer + self.bridging_scale = str(inner_clamp_scale) + if self.bridging_f_inner is not None: + self.bridging_clamp: InnerClamp | None = InnerClamp( + self.bridging_f_inner, self.bridging_f_outer ) self.bridging_switch: BridgingSwitch | None = BridgingSwitch( - self.inner_clamp_r_inner, self.inner_clamp_r_outer + self.bridging_f_inner, self.bridging_f_outer ) else: - self.inner_clamp = None + self.bridging_clamp = None self.bridging_switch = None + # The table sizes a window, so an unbridged descriptor needs none and + # is not asked for the element symbols the covalent scale reads. + self.contact_radius = None + if self.bridging_f_inner is not None: + table = contact_radius_table( + self.ntypes, self.type_map, self.bridging_scale + ) + check_window_inside_cutoff(table, self.bridging_f_outer, self.rcut) + self.contact_radius = table # === Env seed parameters === self.use_env_seed = bool(use_env_seed) @@ -1273,6 +1317,12 @@ def __init__( self.mean = np.zeros(0, dtype=PRECISION_DICT[self.precision]) self.stddev = np.ones(0, dtype=PRECISION_DICT[self.precision]) + def pair_contact( + self, center_type: Array, neighbor_type: Array, dtype: Any + ) -> Array: + """Length scale of each edge in Å, shape (E, 1); see :func:`pair_contact`.""" + return pair_contact(self.contact_radius, center_type, neighbor_type, dtype) + def call( self, coord_ext: Array, @@ -1290,6 +1340,7 @@ def call( Array | None, Array | None, Array | None, + Array | None, ]: """ Compute the descriptor. @@ -1332,6 +1383,12 @@ def call( None (not used). sw None (not used). + node_gate + Per-atom source gate of the bridging window with shape + (nf, nloc, 1), which fades the learned atomic energy; ``None`` + for an unbridged descriptor. DPA4 always returns this sixth + element beyond the common five-element descriptor contract, and + the atomic model reads it by position. Raises ------ @@ -1381,7 +1438,7 @@ def call( graph, atype_flat = _graph_from_padded_nlist( coord_ext, atype_ext, nlist, mapping ) - x_scalar, _ = self._run_graph( + x_scalar, _, node_gate = self._run_graph( graph, atype_flat, n_out_nodes=nf * nloc, @@ -1392,12 +1449,15 @@ def call( # ``_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)) + if node_gate is not None: + node_gate = xp.reshape(node_gate, (nf, nloc, 1)) return ( descriptor, None, None, None, None, + node_gate, ) def call_graph( @@ -1408,7 +1468,7 @@ def call_graph( comm_dict: dict[str, Array] | None = None, spin: Array | None = None, charge_spin: Array | None = None, - ) -> tuple[Array, None]: + ) -> tuple[Array, None, Array | None]: """Graph-native descriptor forward on the flat node axis. Parameters @@ -1447,9 +1507,11 @@ def call_graph( Returns ------- - tuple[Array, None] - Flat ``(N, channels)`` descriptor in global precision, and - ``None`` (DPA4 produces no equivariant rot_mat for the fitting). + tuple[Array, None, Array | None] + Flat ``(N, channels)`` descriptor in global precision, ``None`` + (DPA4 produces no equivariant rot_mat for the fitting), and for + a bridged descriptor the flat ``(N,)`` per-node source gate that + fades the learned atomic energy (``None`` otherwise). Raises ------ @@ -1464,7 +1526,7 @@ def call_graph( nf=nf, ref=graph.edge_vec, ) - x_scalar, _ = self._run_graph( + x_scalar, _, node_gate = self._run_graph( graph, atype, charge_spin=charge_spin, spin=spin, comm_dict=comm_dict ) # ``_run_graph`` returns the read-out with its SO(3) singleton @@ -1472,7 +1534,7 @@ def call_graph( # 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 + return x_scalar, None, node_gate def _run_graph( self, @@ -1484,7 +1546,7 @@ def _run_graph( charge_spin: Array | None = None, spin: Array | None = None, comm_dict: dict[str, Array] | None = None, - ) -> tuple[Array, Array]: + ) -> tuple[Array, Array, Array | None]: """Graph-native descriptor forward shared by both descriptor entries. Both public entries -- the dense-nlist ``call`` adapter and the @@ -1529,10 +1591,13 @@ def _run_graph( Returns ------- - tuple[Array, Array] + tuple[Array, Array, Array | None] 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 + ``(n_out_nodes, 1, 1, channels)``, in global precision, the + full multipole feature tensor ``x``, and the per-node source gate + of the read-out rows with shape ``(n_out_nodes,)`` in global + precision for a bridged descriptor (``None`` otherwise). The + public entries flatten the singleton axes. Raises @@ -1584,16 +1649,30 @@ def _run_graph( ) # === Step 3. Build edge cache once (sparse edges) === training = self._in_training_mode() + compute_dtype = get_xp_precision(xp, self.compute_precision) + # The window of every edge is sized by the two types at its ends, so + # the descriptor, which owns the radius table, resolves it here and the + # edge cache receives one per-edge length scale. + edge_contact = ( + None + if self.contact_radius is None + else self.pair_contact( + xp.take(atype_flat, edge_index[1, ...]), + xp.take(atype_flat, edge_index[0, ...]), + compute_dtype, + ) + ) edge_cache = _edge_cache_from_arrays( type_ebed=type_ebed, edge_index=edge_index, edge_vec=edge_vec, edge_mask=edge_mask, - compute_dtype=get_xp_precision(xp, self.compute_precision), + compute_dtype=compute_dtype, eps=self.eps, deg_norm_floor=(self.deg_norm_floor if self.version >= 1.1 else self.eps), - inner_clamp=self.inner_clamp, + bridging_clamp=self.bridging_clamp, bridging_switch=self.bridging_switch, + edge_contact=edge_contact, edge_envelope=self.edge_envelope, radial_basis=self.radial_basis, fused_radial=None if training else self._cuda_radial_fn, @@ -1738,7 +1817,14 @@ def _run_graph( # === Step 12. Final l=0 output mixing === x_scalar = self._apply_readout(x, n_out_nodes) - return xp.astype(x_scalar, get_xp_precision(xp, "global")), x + node_gate = ( + None + if edge_cache.node_gate is None + else xp.astype( + edge_cache.node_gate[:n_out_nodes], get_xp_precision(xp, "global") + ) + ) + return xp.astype(x_scalar, get_xp_precision(xp, "global")), x, node_gate def _forward_blocks( self, @@ -2772,9 +2858,11 @@ def _migrate_variables( anything is assigned to a module: ``load_state_dict`` restores a module's own buffers before descending into its children, so a migration applied to live attributes would rewrite values the child - load is about to overwrite. Only representations are upgraded here; - a difference no rewrite can absorb stays a forward-time branch on - :attr:`version`, so a migrated descriptor never changes its own math. + load is about to overwrite. Both checkpoints and portable records + pass through this migration. A bridged receiving descriptor rejects + legacy states before its tensors are restored. Other differences that + no rewrite can absorb stay forward-time branches on :attr:`version`, + so a migrated descriptor never changes its own math. Version 1.2 moved the env-seed spin gate from the spin coordinate to the resulting environment quadratic form. For an active-spin model, @@ -2785,6 +2873,14 @@ def _migrate_variables( be activated by fine-tuning. Versions below 1.1 predate the native-spin route and retain their original forward semantics. + Version 1.3 changes only bridging. A version-1.2 state loaded into an + unbridged descriptor adopts the current record version without changing + its function. A bridged receiving descriptor rejects every pre-1.3 state: + state dictionaries do not identify the source's bridging configuration. + Fine-tuning first loads a compatible plain checkpoint into its original + unbridged model, migrates it, then transfers the state to the bridged + target. Legacy bridged checkpoints cannot be migrated this way. + Parameters ---------- variables @@ -2799,28 +2895,33 @@ def _migrate_variables( float Version the variables express after migration. """ - if not 1.1 <= version < 1.2: - return version - - gate_key = prefix + "env_seed_embedding.spin_scale" - if self.use_spin is not None and not any(self.use_spin): - # dpmodel serialization names NativeLayer weights ``matrix``; - # pt_expt state dictionaries expose the wrapped attribute as ``w``. - dormant_keys = ( - "spin_embedding.mag_layer2.matrix", - "spin_embedding.mag_layer2.w", - "spin_embedding.adam_spin_vec_weight", - "spin_embedding.adam_spin_nbr_weight", - "env_seed_embedding.spin_scale", - ) - for name in dormant_keys: - key = prefix + name - if key in variables: - xp = array_api_compat.array_namespace(variables[key]) - variables[key] = xp.zeros_like(variables[key]) - elif gate_key in variables: - variables[gate_key] = variables[gate_key] ** 2 - return 1.2 + check_bridging_record_version( + {"inner_clamp_f_inner": self.bridging_f_inner}, version + ) + if 1.1 <= version < 1.2: + gate_key = prefix + "env_seed_embedding.spin_scale" + if self.use_spin is not None and not any(self.use_spin): + # dpmodel serialization names NativeLayer weights ``matrix``; + # pt_expt state dictionaries expose the wrapped attribute as ``w``. + dormant_keys = ( + "spin_embedding.mag_layer2.matrix", + "spin_embedding.mag_layer2.w", + "spin_embedding.adam_spin_vec_weight", + "spin_embedding.adam_spin_nbr_weight", + "env_seed_embedding.spin_scale", + ) + for name in dormant_keys: + key = prefix + name + if key in variables: + xp = array_api_compat.array_namespace(variables[key]) + variables[key] = xp.zeros_like(variables[key]) + elif gate_key in variables: + variables[gate_key] = variables[gate_key] ** 2 + version = 1.2 + + if 1.2 <= version < BRIDGING_RECORD_VERSION: + return BRIDGING_RECORD_VERSION + return version def serialize(self) -> dict[str, Any]: return { @@ -2889,8 +2990,9 @@ def serialize(self) -> dict[str, Any]: "eps": self.eps, "trainable": self.trainable, "seed": self.seed, - "inner_clamp_r_inner": self.inner_clamp_r_inner, - "inner_clamp_r_outer": self.inner_clamp_r_outer, + "inner_clamp_f_inner": self.bridging_f_inner, + "inner_clamp_f_outer": self.bridging_f_outer, + "inner_clamp_scale": self.bridging_scale, "add_chg_spin_ebd": self.add_chg_spin_ebd, "default_chg_spin": self.default_chg_spin, "use_spin": self.use_spin, @@ -2914,6 +3016,7 @@ def deserialize(cls, data: dict[str, Any]) -> DescrptDPA4: variables = data.pop("@variables") data.pop("env_mat", None) config.pop("s2_grid_resolution", None) + migrate_inner_clamp_keys(config) obj = cls(**config) obj.version = obj._migrate_variables(variables, version) obj._load_variables(variables) diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/__init__.py b/deepmd/dpmodel/descriptor/dpa4_nn/__init__.py index 817b50c841..b0f471ec33 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/__init__.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/__init__.py @@ -31,7 +31,7 @@ from .edge_cache import ( EdgeCache, build_edge_type_feat, - compute_edge_src_gate, + compute_source_gates, edge_cache_to_dtype, ) from .embedding import ( @@ -91,11 +91,13 @@ resolve_so3_grid, ) from .radial import ( + BridgingClamp, BridgingSwitch, C3CutoffEnvelope, InnerClamp, RadialBasis, RadialMLP, + pair_contact, ) from .so2 import ( DynamicRadialDegreeMixer, @@ -127,6 +129,7 @@ "ATTN_RES_MODES", "BaseGridNet", "BaseGridProjector", + "BridgingClamp", "BridgingSwitch", "C3CutoffEnvelope", "ChannelLinear", @@ -176,7 +179,7 @@ "build_m_major_l_index", "build_merged_state_dict", "build_rotate_inv_rescale", - "compute_edge_src_gate", + "compute_source_gates", "edge_cache_to_dtype", "fold_lora_state_dict_keys", "get_promoted_dtype", @@ -185,6 +188,7 @@ "init_trunc_normal_fan_in_out", "map_degree_idx", "merge_lora_into_base", + "pair_contact", "project_D_to_m", "project_Dt_from_m", "quaternion_multiply", diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/attention.py b/deepmd/dpmodel/descriptor/dpa4_nn/attention.py index a06e9fe2be..d7a9bef0d4 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/attention.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/attention.py @@ -28,17 +28,6 @@ ) -def _stop_gradient(value: Any) -> Any: - """Return ``value`` with backend gradient tracking disabled.""" - if array_api_compat.is_torch_array(value): - return value.detach() - if array_api_compat.is_jax_array(value): - import jax - - return jax.lax.stop_gradient(value) - return value - - def segment_envelope_gated_softmax( logits: Any, edge_env: Any, @@ -46,12 +35,18 @@ def segment_envelope_gated_softmax( n_nodes: int, z_bias_raw: Any, eps: float, - src_weight: Any = None, edge_mask: Any = None, ) -> Any: """ Compute destination-wise envelope-gated softmax attention. + The physical mass of an edge is ``edge_env**2 * exp(logits)`` and the + denominator of its destination is the sum of the edge masses plus the + positive null mass ``softplus(z_bias_raw) + eps``. An edge whose envelope + is zero is therefore absent from both, which is what the frozen-zone + invariance of a bridged model needs, since its envelope carries the + source gate. + Parameters ---------- logits @@ -72,16 +67,6 @@ def segment_envelope_gated_softmax( Softplus is applied to keep the bias strictly positive. eps Small positive floor added to the physical null mass. - src_weight - Optional per-edge source-side multiplier with shape (E, 1) or - (E,). When provided, the physical per-edge mass is - ``edge_env**2 * src_weight * exp(logits)`` and the denominator is the - sum of edge masses plus the positive null mass - ``softplus(z_bias_raw) + eps``. - ``src_weight = 0`` therefore removes the source from both the - numerator and the denominator, which is what SFPG needs so that - a muted source does not even leak through the softmax - normalization. edge_mask Optional binary padded-edge validity mask with shape (E,) or (E, 1); one marks valid slots and zero marks invalid slots. Invalid slots drop @@ -101,28 +86,15 @@ def segment_envelope_gated_softmax( compute_dtype = xp.float32 if promote else input_dtype dst = xp.astype(dst, xp.int64) - # === Step 1. Build factor-wise effective logits === - # Computing the logarithms before multiplying the factors avoids losing a - # physically nonzero edge when ``edge_env**2 * src_weight`` underflows. + # === Step 1. Build the effective logits === + # Taking the logarithm of the envelope before adding it to the logits + # avoids losing a physically nonzero edge when ``edge_env**2`` underflows. logits_2d = xp.astype(xp.reshape(logits, (n_edge, n_channel)), compute_dtype) edge_env_1d = xp.astype(xp.reshape(edge_env, (n_edge,)), compute_dtype) edge_positive = edge_env_1d > 0.0 ones = xp.ones((n_edge,), dtype=compute_dtype, device=device) log_weight = 2.0 * xp.log(xp.where(edge_positive, edge_env_1d, ones)) active = edge_positive - source_ratio = None - if src_weight is not None: - source_weight = xp.astype(xp.reshape(src_weight, (n_edge,)), compute_dtype) - source_positive = source_weight > 0.0 - safe_source = xp.where(source_positive, source_weight, ones) - source_scale = _stop_gradient(safe_source) - log_weight = log_weight + xp.log(source_scale) - source_ratio = xp.where( - source_positive, - source_weight / source_scale, - xp.zeros((n_edge,), dtype=compute_dtype, device=device), - ) - active = active & source_positive if edge_mask is not None: mask = xp.astype(xp.reshape(edge_mask, (n_edge,)), compute_dtype) mask_positive = mask > 0.0 @@ -162,11 +134,6 @@ def segment_envelope_gated_softmax( # === Step 3. Normalize edge and null masses in the shared shifted frame === edge_exp = xp.exp(effective_logits - edge_max) - if source_ratio is not None: - # ``source_scale`` carries the forward magnitude in log space, while - # this linear ratio carries derivatives without differentiating - # ``log(src_weight)`` at extremely small positive gates. - edge_exp = edge_exp * source_ratio[:, None] denom_sum = xp_add_at( xp.zeros((n_nodes, n_channel), dtype=compute_dtype, device=device), dst, diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py b/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py index ec51a3d23e..a104c8e16e 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py @@ -70,7 +70,20 @@ class EdgeCache: Radial basis with shape (E, n_radial). The C^3 cutoff envelope is already baked in. edge_env - C^3 cutoff envelope weights with shape (E, 1). + C^3 cutoff envelope weights with shape (E, 1). In bridging mode the + Source Freeze Propagation Gate is folded into it: the envelope of the + edge ``j -> k`` is multiplied by the product of the switching + amplitudes ``w(r_{jl})`` of every pair of the source ``j`` other than + the pair ``(j, k)`` itself, so an edge whose source has a third + neighbor in the frozen zone vanishes everywhere the envelope enters + (messages, environment seed, attention masses and degree), exactly as + an edge beyond the cutoff does, while the edge between the two atoms + of a frozen pair keeps its clamped geometry. + node_gate + Per-node source gate ``eta[j] = prod_{k in N(j)} w(r_{jk})`` with + shape (N,), the product over every pair of the node; ``None`` without + bridging. The atomic model fades the learned atomic energy of node + ``j`` with it. deg Envelope-squared smooth degree with shape (N,), computed as ``sum(edge_env**2)`` over incoming edges. @@ -95,17 +108,6 @@ class EdgeCache: Lazy cache for endpoint CSR views used by segmented accelerated operators, keyed by endpoint role (``"dst"`` or ``"src"``). Built once per step and shared by every consumer. - edge_src_gate - Optional per-edge Source Freeze Propagation Gate (SFPG) weight with - shape (E, 1). Equals ``eta[src]`` where - ``eta[j] = prod_{k in N(j)} w(r_{jk})`` and ``w`` is the - :class:`BridgingSwitch` C3 switching amplitude. Present only when - the model runs in bridging mode; ``None`` otherwise. Aggregation - sites (``GeometricInitialEmbedding``, ``EnvironmentInitialEmbedding``, - ``SO2Convolution``) multiply their per-edge message contribution - by this gate to forbid any node whose local neighborhood enters - the frozen zone from propagating information along its outgoing - edges. edge_mask 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 @@ -127,73 +129,81 @@ class EdgeCache: D_to_m_cache: dict[str, Any] = field(default_factory=dict) Dt_from_m_cache: dict[str, Any] = field(default_factory=dict) csr_cache: dict[str, Any] | None = field(default_factory=dict) - edge_src_gate: Any = None edge_quat: Any = None edge_mask: Any = None + node_gate: Any = None -def compute_edge_src_gate( +def compute_source_gates( *, edge_len: Any, + edge_contact: Any, src: Any, n_nodes: int, - bridging_switch: Callable[[Any], Any], + bridging_switch: Callable[[Any, Any], Any], edge_keep_f: Any = None, node_partial_exchange: Callable[[Any], Any] | None = None, -) -> Any: +) -> tuple[Any, Any]: """ - Compute the per-edge source gate for SFPG from edge lengths. - - The gate implements a per-node "non-frozen confidence" and broadcasts - it back to edges along the source axis:: - - w_e = bridging_switch(edge_len_e) in [0, 1] - eta_j = prod_{e: src_e = j} w_e in [0, 1] - gate_e = eta_{src_e} in [0, 1] - - ``w_e = 0`` at ``r_{jk} <= r_inner`` ensures ``eta_j = 0`` for any - node with at least one neighbor in the frozen zone. Masked edges - (padding, excluded type pairs) must contribute the multiplicative - identity ``1`` so they never spuriously mute a valid source node; - callers supply ``edge_keep_f`` for this. - - The product is **not** realised by ``scatter_reduce(reduce="prod")``: - its registered backward handles exact zeros with a data-dependent - "count leave-one-out" branch that creates unbacked symints under - ``make_fx(tracing_mode="symbolic")`` and breaks the SeZM compile - path's double-backward tracing. Instead, the product is decomposed - into a log-sum on non-zero contributions combined with an explicit - "any zero per group" indicator that routes the frozen case through - ``where``. Both branches use only shape-preserving standard - ops (``scatter_add``, ``where``, ``exp``, ``log``) with backed - symints, so the graph survives symbolic tracing cleanly. + Compute the per-node and per-edge source gates of SFPG from edge lengths. + + The per-node gate is the "non-frozen confidence" of a node, the per-edge + gate the same product with the edge's own pair left out:: + + w_e = bridging_switch(edge_len_e, edge_contact_e) in [0, 1] + eta_j = prod_{e: src_e = j} w_e in [0, 1] + gate_e = eta_{src_e} / w_e = prod_{e': src_e' = src_e, e' != e} w_e' + + ``w_e = 0`` inside the inner radius of that edge's own pair ensures + ``eta_j = 0`` for any node with at least one neighbor in the frozen zone, + and ``gate_e = 0`` for every edge such a node emits except the edge to + that frozen neighbor. The cache builder multiplies ``gate_e`` into the + edge envelope, which every consumer weights by, so a muted edge is + indistinguishable from a deleted one: the frozen atom disappears from the + environment of every third atom, while the two atoms of a frozen pair + keep seeing each other at the clamped distance. The fitting scales the + deviation of its output from the bias of the type by ``eta_j``, so the + frozen atom contributes that bias alone. Masked edges (padding, + excluded type pairs) must contribute the multiplicative identity ``1`` so + they never spuriously mute a valid source node; callers supply + ``edge_keep_f`` for this. + + The product is decomposed into a log-sum on non-zero contributions + combined with an explicit "any zero per group" indicator that routes the + frozen case through ``where``; the leave-one-out product subtracts the + edge's own log-amplitude and zero flag from the totals of its source. + Both branches use only shape-preserving standard ops, mirroring the pt + implementation so the two backends agree bitwise in float64. The gradient consequence at the plateau is exact: ``BridgingSwitch`` - places ``w'(r) = 0`` for every ``r <= r_inner``, so the chain rule + places ``w'(r) = 0`` everywhere inside the inner radius, so the chain rule ``d eta / d r = (leave-one-out factor) * w'(r) = anything * 0 = 0`` - holds regardless of how the muted ``where`` branch treats the - upstream gradient. In the transition zone every edge has strictly - positive ``w`` and the log-sum branch gives the standard product - gradient. + holds regardless of how the muted ``where`` branch treats the upstream + gradient. In the transition zone every edge has strictly positive ``w`` + and the log-sum branch gives the standard product gradient. Parameters ---------- edge_len Per-edge distances with shape (E, 1). + edge_contact + Per-edge length scale with shape (E, 1) in Å, the sum of the two + endpoint radii. The switch measures the window against it, so every + element combination gets its own pair of radii. src Source node indices with shape (E,). n_nodes Total number of nodes N. bridging_switch - Callable ``r -> w(r)`` with ``w: [0, ∞) -> [0, 1]``, typically a - :class:`BridgingSwitch` instance. + Callable ``(r, contact) -> w`` with ``w: [0, ∞) -> [0, 1]``, typically + a :class:`BridgingSwitch` instance. edge_keep_f Optional per-edge keep weights with shape (E, 1), with ``0`` on masked edges and ``1`` on kept edges. If provided, masked edges are rewritten to ``w = 1`` before the product reduction. node_partial_exchange Optional cross-rank completion hook for the per-node partials - (issue #5906). Receives the ``(n_nodes, 2)`` float tensor + (issue #5906). Receives the ``(n_nodes, 2)`` float array ``[log_eta, zero_count]`` and returns the globally completed one (reverse-accumulate ghost rows into owners, then broadcast the completed owner values back onto ghosts). ``None`` (the default) @@ -202,14 +212,16 @@ def compute_edge_src_gate( Returns ------- - Array - Per-edge source gate with shape (E, 1), aligned on the same edge - axis as the rest of the cache. + node_gate : Array + Per-node source gate ``eta`` with shape (N,). + edge_gate : Array + Per-edge leave-one-out gate with shape (E, 1), aligned on the same + edge axis as the rest of the cache. """ xp = array_api_compat.array_namespace(edge_len, src) device = array_api_compat.device(edge_len) # === Step 1. Per-edge switching amplitude w(r) in [0, 1] === - edge_w = bridging_switch(edge_len) # (E, 1) + edge_w = bridging_switch(edge_len, edge_contact) # (E, 1) if edge_keep_f is not None: # Force w = 1 on masked edges so they are neutral for the product. edge_w = edge_w * edge_keep_f + (1.0 - edge_keep_f) @@ -233,10 +245,11 @@ def compute_edge_src_gate( # the hard-freeze rule. Float count (values are small integers, exact # in fp) so both partials ride ONE border-exchange tensor when # completing across ranks. + is_zero_f = xp.astype(is_zero, edge_w.dtype) zero_count = xp_add_at( xp.zeros((n_nodes,), dtype=edge_w.dtype, device=device), src, - xp.astype(is_zero, edge_w.dtype), + is_zero_f, ) # === Step 3b. Cross-rank completion of the per-node partials === @@ -251,12 +264,17 @@ def compute_edge_src_gate( log_eta = packed[..., 0] zero_count = packed[..., 1] - eta_nonzero_path = xp.exp(log_eta) - any_zero = zero_count > 0.5 + # === Step 4. Per-node gate === + node_gate = xp.where(zero_count > 0.5, xp.zeros_like(log_eta), xp.exp(log_eta)) - # === Step 4. Combine and broadcast back to edges via source === - eta = xp.where(any_zero, xp.zeros_like(eta_nonzero_path), eta_nonzero_path) - return xp.take(eta, src, axis=0)[:, None] + # === Step 5. Per-edge leave-one-out gate === + # The edge's own factor is removed from its source's totals: its + # log-amplitude (0 when it is a frozen edge, because ``safe_w`` is 1 there) + # and its zero flag, so the remaining product is over the other pairs. + loo_log = xp.take(log_eta, src, axis=0) - log_safe # (E,) + loo_zero = xp.take(zero_count, src, axis=0) - is_zero_f # (E,) + edge_gate = xp.where(loo_zero > 0.5, xp.zeros_like(loo_log), xp.exp(loo_log)) + return node_gate, edge_gate[:, None] def _edge_cache_from_arrays( @@ -268,8 +286,9 @@ def _edge_cache_from_arrays( compute_dtype: Any, eps: float, deg_norm_floor: float, - inner_clamp: Callable[[Any], Any] | None, - bridging_switch: Callable[[Any], Any] | None, + bridging_clamp: Callable[[Any, Any], Any] | None, + bridging_switch: Callable[[Any, Any], Any] | None, + edge_contact: Any, edge_envelope: Callable[[Any], Any], radial_basis: Callable[[Any], Any], random_gamma: bool, @@ -306,15 +325,22 @@ def _edge_cache_from_arrays( deg_norm_floor Floor added to the envelope-squared degree before inverse-sqrt normalization (see :func:`_finalize_edge_cache`). - inner_clamp - Optional inner clamp used to freeze short-range geometry below `r_inner`. + bridging_clamp + Optional distance clamp that freezes the geometry the descriptor sees + for a pair inside its bridging window (see :class:`InnerClamp`). bridging_switch - Optional C3 switching amplitude ``w(r) -> [0, 1]`` that drives - the Source Freeze Propagation Gate. When provided, a per-edge - ``edge_src_gate`` is computed from the node-wise product of - ``w(r_{jk})`` along each source node's outgoing edges. Masked + Optional C3 switching amplitude ``w(r, contact) -> [0, 1]`` that drives + the Source Freeze Propagation Gate. When provided, the product of + ``w(r_{jl})`` over the true lengths of every pair of the source node + other than the edge's own pair is folded into the envelope of every + edge that node emits, and the full product is stored on the cache as + ``node_gate``. Masked edges (``edge_keep=False``) are forced to ``w=1`` so they never leak into the product. + edge_contact + Per-edge length scale with shape (E, 1) in Å that both window modules + measure their fractions against, or ``None`` when neither is given. + The descriptor owns the per-type radii and resolves it before the call. edge_envelope C^3 edge envelope module. radial_basis @@ -338,7 +364,7 @@ def _edge_cache_from_arrays( ``[0, 2*pi)``; callers may inject angles to pin a draw. node_partial_exchange Optional cross-rank completion hook forwarded to - :func:`compute_edge_src_gate` (issue #5906); only meaningful when + :func:`compute_source_gates` (issue #5906); only meaningful when ``bridging_switch`` is provided. Returns @@ -367,8 +393,9 @@ def _edge_cache_from_arrays( # === Step 3. Edge length, envelope, and radial basis === edge_len = safe_norm(edge_vec, eps) - if inner_clamp is not None: - clamped = inner_clamp(edge_len) + edge_len_true = edge_len + if bridging_clamp is not None: + clamped = bridging_clamp(edge_len, edge_contact) scale = clamped / edge_len edge_vec = edge_vec * scale edge_len = clamped @@ -394,20 +421,29 @@ def _edge_cache_from_arrays( edge_type_feat = edge_type_feat * xp.astype(edge_keep_f, edge_type_feat.dtype) # === Step 6. Source Freeze Propagation Gate (optional) === - # The sparse-edge path packs masked dummy edges so the compiled graph sees - # a statically non-empty, non-singular edge tensor. ``edge_keep_f`` rewrites - # any such slot to ``w=1`` inside ``compute_edge_src_gate``, keeping the - # product reduction unaffected by padding. - edge_src_gate: Any = None + # The leave-one-out gate is folded into the envelope, the one factor every + # consumer of an edge weights by: the messages, the environment seed, the + # attention masses and the degree normalization all see a muted edge + # exactly as they see an edge beyond the cutoff. The gate reads the true + # pair length, so it closes where the clamp above has frozen the geometry. + # The radial basis stays ungated because it feeds networks whose output + # the envelope multiplies. The sparse-edge path packs masked dummy edges + # so the compiled graph sees a statically non-empty, non-singular edge + # tensor; ``edge_keep_f`` rewrites any such slot to ``w=1`` inside + # ``compute_source_gates``, keeping the product reduction unaffected by + # padding. The per-node product travels on the cache to the atomic model. + node_gate = None if bridging_switch is not None: - edge_src_gate = compute_edge_src_gate( - edge_len=edge_len, + node_gate, edge_gate = compute_source_gates( + edge_len=edge_len_true, + edge_contact=edge_contact, src=src, n_nodes=n_nodes, bridging_switch=bridging_switch, edge_keep_f=edge_keep_f, node_partial_exchange=node_partial_exchange, ) + edge_env = edge_env * edge_gate return _finalize_edge_cache( n_nodes=n_nodes, @@ -421,7 +457,7 @@ def _edge_cache_from_arrays( Dt_full=Dt_full, edge_quat=edge_quat, deg_norm_floor=deg_norm_floor, - edge_src_gate=edge_src_gate, + node_gate=node_gate, ) @@ -509,7 +545,7 @@ def _finalize_edge_cache( Dt_full: Any, edge_quat: Any, deg_norm_floor: float, - edge_src_gate: Any = None, + node_gate: Any = None, ) -> EdgeCache: """ Assemble the shared `EdgeCache` layout. @@ -544,9 +580,9 @@ def _finalize_edge_cache( normalization. A tiny ``eps`` reproduces the legacy behavior; an ``O(1)`` value makes sparse-neighborhood features vanish smoothly at ``rcut`` instead of saturating and kinking. - edge_src_gate - Optional per-edge SFPG weight with shape (E, 1). ``None`` in - non-bridging mode. + node_gate + Per-node source gate with shape (N,) of a bridged descriptor, or + ``None``. Returns ------- @@ -577,8 +613,8 @@ def _finalize_edge_cache( D_to_m_cache={}, Dt_from_m_cache={}, csr_cache={}, - edge_src_gate=edge_src_gate, edge_quat=edge_quat, + node_gate=node_gate, ) @@ -619,8 +655,9 @@ def edge_cache_to_dtype(cache: EdgeCache, dtype: Any) -> EdgeCache: """ Convert all floating-point tensors in EdgeCache to the specified dtype. - Integer tensors (src, dst) are unchanged. This is a standalone function - (not a method) to keep it side-effect free. + Integer tensors (src, dst) are unchanged, and so is the node gate, which + no block reads: it is applied to the fitting output in global precision. + This is a standalone function (not a method) to keep it side-effect free. Parameters ---------- @@ -639,18 +676,14 @@ def edge_cache_to_dtype(cache: EdgeCache, dtype: Any) -> EdgeCache: # Use local variables with explicit None check and assignment. _D_full = cache.D_full _Dt_full = cache.Dt_full - _edge_src_gate = cache.edge_src_gate _edge_quat = cache.edge_quat D_full: Any = None Dt_full: Any = None - edge_src_gate: Any = None edge_quat: Any = None if _D_full is not None: D_full = xp.astype(_D_full, dtype) if _Dt_full is not None: Dt_full = xp.astype(_Dt_full, dtype) - if _edge_src_gate is not None: - edge_src_gate = xp.astype(_edge_src_gate, dtype) if _edge_quat is not None: edge_quat = xp.astype(_edge_quat, dtype) @@ -670,7 +703,7 @@ def edge_cache_to_dtype(cache: EdgeCache, dtype: Any) -> EdgeCache: D_to_m_cache=None if cache.D_to_m_cache is None else {}, Dt_from_m_cache=None if cache.Dt_from_m_cache is None else {}, csr_cache=None if cache.csr_cache is None else dict(cache.csr_cache), - edge_src_gate=edge_src_gate, edge_quat=edge_quat, edge_mask=cache.edge_mask, + node_gate=cache.node_gate, ) diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py index 9de135e641..b070cba778 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py @@ -280,8 +280,8 @@ def call( Optional per-edge neighbor-spin l=1 message with shape (E, 3, C) for the native spin scheme (built by ``SpinEmbedding.edge_l1``). It is added to the l=1 rows of the per-edge message, so it shares this - module's source gate, scatter and degree normalization with the - geometric message. + module's scatter and degree normalization with the geometric + message. Returns ------- @@ -317,12 +317,11 @@ def call( # === Step 3. Broadcast radial features per row === # Each non-scalar packed row reuses the radial feature of its degree l. - # The fused operator spans this broadcast and the scatter of Step 5, so + # The fused operator spans this broadcast and the scatter of Step 4, so # it takes over whenever nothing else joins the message in between. if ( self._can_fuse_scatter(zonal_coupling) and spin_l1_message is None - and edge_cache.edge_src_gate is None and edge_cache.csr_cache is not None ): return self.forward_fused_scatter( @@ -342,7 +341,7 @@ def call( # === Step 3b. Fold in the neighbor-spin l=1 message (native spin) === # The l=1 coefficients occupy the first three packed non-scalar rows, so # the neighbor-spin message joins the geometric message there and then - # shares the source gate, scatter and degree normalization below. + # shares the scatter and degree normalization below. if spin_l1_message is not None: l1_local_index = xp_asarray_nodetach(xp, self.l1_local_index, device=device) scatter_index = xp.broadcast_to( @@ -352,17 +351,7 @@ def call( non_scalar_message, 1, scatter_index, spin_l1_message ) - # === Step 4. Source Freeze Propagation Gate (optional) === - # Mute messages emitted by nodes whose local neighborhood enters - # the frozen zone. ``edge_src_gate`` is ``None`` outside bridging - # mode so this is a no-op in normal training. - src_gate = edge_cache.edge_src_gate - if src_gate is not None: - non_scalar_message = non_scalar_message * xp.astype( - xp.reshape(src_gate, (n_edge, 1, 1)), non_scalar_message.dtype - ) - - # === Step 5. Scatter to nodes and normalize === + # === Step 4. Scatter to nodes and normalize === # Destination scatter-add over ``edge_cache.dst`` (pt ``index_add_``), # applied after the validity masking below. This reduction is # layout-agnostic: it is correct both for the padded ``call`` (row-major @@ -729,13 +718,6 @@ def call( outer_flat = xp.reshape( outer, (n_edge, self.coord_dim * self.embed_dim) ) # (E, coord_dim*embed_dim) - # Source Freeze Propagation Gate: mute the outer-product contribution - # of any edge whose source node has a neighbor in the frozen zone. - src_gate = edge_cache.edge_src_gate - if src_gate is not None: - outer_flat = outer_flat * xp.astype( - xp.reshape(src_gate, (n_edge, 1)), outer_flat.dtype - ) # 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 @@ -1200,9 +1182,9 @@ def edge_l1( (neighbor) spin, scaled by a per-source-type per-channel weight and gated by the C^3 envelope. The message is returned per edge; the geometric initial embedding folds it into the l=1 rows and applies the - shared source gate, scatter and degree normalization, so a neighbor's - spin direction enters an atom's l=1 backbone before any interaction - block (the spin analogue of the geometric initial embedding). + shared scatter and degree normalization, so a neighbor's spin direction + enters an atom's l=1 backbone before any interaction block (the spin + analogue of the geometric initial embedding). Parameters ---------- diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/radial.py b/deepmd/dpmodel/descriptor/dpa4_nn/radial.py index 3c06cc2f9f..183ce718d8 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/radial.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/radial.py @@ -39,6 +39,9 @@ from deepmd.dpmodel.utils.seed import ( child_seed, ) +from deepmd.utils.bridging import ( + window_midpoint, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -269,65 +272,125 @@ def call(self, dst: Any) -> Any: return u**4 * series -class InnerClamp(NativeOP): +def pair_contact( + contact_radius: Any, center_type: Any, neighbor_type: Any, dtype: Any +) -> Any: """ - C3-continuous inner distance clamping for zone bridging. + Gather the length scale of each edge from the types at its ends. - Applies a septic Hermite polynomial transition that freezes distances - below ``r_inner`` to the constant ``r_inner``, then smoothly transitions - back to identity at ``r_outer``:: + Parameters + ---------- + contact_radius : Any + Per-type length scales in Å with shape (ntypes + 1,), the trailing + entry belonging to the padding type. + center_type : Any + Type index of the center atom of every edge, with shape (E,). + neighbor_type : Any + Type index of the neighbor atom of every edge, with shape (E,). + dtype : Any + Floating-point dtype of the result. The edge distances are divided by + that result, so this is the dtype the geometry is reduced in rather + than the dtype of the incoming coordinates. - r̃(r) = r_inner if r <= r_inner - r̃(r) = r_inner + (r_outer - r_inner) * h(t) if r_inner < r < r_outer - r̃(r) = r if r >= r_outer + Returns + ------- + Any + Contact distance of each edge in Å, with shape (E, 1). + """ + xp = array_api_compat.array_namespace(center_type) + radius = xp.asarray( + contact_radius, dtype=dtype, device=array_api_compat.device(center_type) + ) + contact = xp.take(radius, center_type) + xp.take(radius, neighbor_type) + return contact[:, None] - h(t) = 20t^4 - 45t^5 + 36t^6 - 10t^7, t = (r - r_inner) / (r_outer - r_inner) - Boundary conditions: - ``h(0)=0``, ``h(1)=1``, ``h'(0)=0``, ``h'(1)=1``, - ``h''(0)=0``, ``h''(1)=0``, ``h'''(0)=0``, ``h'''(1)=0``. - This ensures C3 continuity: ``dr̃/dr = 0`` at r_inner (frozen zone) and - ``dr̃/dr = 1`` at r_outer (identity zone), with matched second and third - derivatives at both boundaries. +class InnerClamp(NativeOP): + """ + C3-continuous distance clamp of the DPA4 bridging window. + + The window of a pair is two dimensionless fractions of that pair's own + length scale ``s``, the sum of the covalent radii of its two elements, so + one window shape serves every element combination. In the reduced distance + ``u = r / s`` the clamp freezes everything below the freeze point + ``f_z`` at that value and returns to the identity at ``f_outer``:: + + ũ = f_z if u <= f_z + ũ = f_z + (f_outer - f_z) * h(t) if f_z < u < f_outer + ũ = u if u >= f_outer + + h(t) = 20t^4 - 45t^5 + 36t^6 - 10t^7, t = (u - f_z) / (f_outer - f_z) + f_z = f_inner + 0.4 (f_outer - f_inner) + + and the clamped distance is ``r̃ = s * ũ``. Boundary conditions + ``h(0)=0``, ``h(1)=1``, ``h'(0)=0``, ``h'(1)=1``, ``h''(0)=h''(1)=0``, + ``h'''(0)=h'''(1)=0`` give C3 continuity: ``dr̃/dr = 0`` at the freeze + point and ``dr̃/dr = 1`` at the outer radius (identity zone), with + matched second and third derivatives at both. + + The freeze point sits two fifths into the window, one tenth of the width + below the frame-filter point at the midpoint, where the training data of a + bridged model end. The displayed distance therefore still moves at the + filter point, with slope ``h'(1/6) = 0.22``, so the network can fit the + labels just above it through the pair's own distance; frozen at the filter + point itself, the learned pair energy could only follow the switch + amplitude, whose slope at the midpoint pins the wall to a stiffness the + labels do not have. Below the filter point the displayed distance keeps + falling by ``0.6 h(1/6) = 0.62 %`` of the window width (5 mÅ for a C-C + pair) before it freezes; that band is the only range of displayed + distances no retained frame covers, and its shallowness is what bounds the + network's extrapolation there. The switch of the window + (:class:`BridgingSwitch`) spans the whole window, so below the freeze point + the learned energy of the pair follows the switch amplitude alone. Parameters ---------- - r_inner : float - Freeze radius in Å. Distances below this are clamped to ``r_inner``. - r_outer : float - Outer boundary of the transition zone in Å. Above this, ``r̃ = r``. + f_inner : float + Inner radius of the window as a fraction of the pair length scale. + f_outer : float + Outer radius as a fraction of the pair length scale. At or above it the + clamp is the identity. Raises ------ ValueError - If ``r_inner >= r_outer`` or either is non-positive. + If ``f_inner >= f_outer`` or either is non-positive. """ - def __init__(self, r_inner: float, r_outer: float) -> None: - if r_inner <= 0 or r_outer <= 0: - raise ValueError("r_inner and r_outer must be positive") - if r_inner >= r_outer: - raise ValueError(f"r_inner ({r_inner}) must be < r_outer ({r_outer})") - self.r_inner = float(r_inner) - self.r_outer = float(r_outer) + #: Position of the freeze point inside the window, as a fraction of its width. + FREEZE_FRACTION = 0.4 + + def __init__(self, f_inner: float, f_outer: float) -> None: + if f_inner <= 0 or f_outer <= 0: + raise ValueError("f_inner and f_outer must be positive") + if f_inner >= f_outer: + raise ValueError(f"f_inner ({f_inner}) must be < f_outer ({f_outer})") + self.f_inner = float(f_inner) + self.f_outer = float(f_outer) + self.f_freeze = self.f_inner + self.FREEZE_FRACTION * ( + self.f_outer - self.f_inner + ) - def call(self, r: Any) -> Any: + def call(self, r: Any, contact: Any) -> Any: """ - Apply inner distance clamping. + Apply the distance clamp. Parameters ---------- r : Array Pair distances with shape (...) or (..., 1) in Å. + contact : Array + Length scale of each pair in Å, broadcastable against ``r``. Returns ------- Array - Clamped distances r̃ with the same shape as input. + Clamped distances r̃ with the same shape as ``r``. """ xp = array_api_compat.array_namespace(r) + u = r / contact t = xp.clip( - (r - self.r_inner) / (self.r_outer - self.r_inner), min=0.0, max=1.0 + (u - self.f_freeze) / (self.f_outer - self.f_freeze), min=0.0, max=1.0 ) t2 = t * t t4 = t2 * t2 @@ -338,11 +401,11 @@ def call(self, r: Any) -> Any: # h''(0)=0, h''(1)=0 # h'''(0)=0, h'''(1)=0 h = t4 * (20.0 + t * (-45.0 + t * (36.0 - 10.0 * t))) - interpolated = self.r_inner + (self.r_outer - self.r_inner) * h - # Identity zone: r >= r_outer returns r directly. - # Both branches have matching first three derivatives at r_outer, + interpolated = contact * (self.f_freeze + (self.f_outer - self.f_freeze) * h) + # Identity zone: u >= f_outer returns r directly. + # Both branches have matching first three derivatives there, # so xp.where preserves C3 continuity here. - return xp.where(r >= self.r_outer, r, interpolated) + return xp.where(u >= self.f_outer, r, interpolated) class BridgingSwitch(NativeOP): @@ -353,11 +416,16 @@ class BridgingSwitch(NativeOP): that measures how far an edge sits outside the frozen zone. It is the elementary piece the Source Freeze Propagation Gate (SFPG) aggregates into a per-node "non-frozen confidence" via a product - over each source node's outgoing edges:: + over each source node's outgoing edges. + + The window of a pair is two dimensionless fractions of that pair's own + length scale ``s``, the sum of the covalent radii of its two elements, so + one window shape serves every element combination. In the reduced distance + ``u = r / s``:: - w(r) = 0 if r <= r_inner (frozen) - w(r) = h((r - r_inner) / (r_outer - r_inner)) if r_inner < r < r_outer (transition) - w(r) = 1 if r >= r_outer (normal) + w = 0 if u <= f_inner (frozen) + w = h((u - f_inner) / (f_outer - f_inner)) if f_inner < u < f_outer (transition) + w = 1 if u >= f_outer (normal) h(t) = 35 t^4 - 84 t^5 + 70 t^6 - 20 t^7 @@ -367,37 +435,39 @@ class BridgingSwitch(NativeOP): h(1)=1, h'(1) = h''(1) = h'''(1) = 0 The vanishing first three derivatives at both endpoints give - ``w \in C^3(\mathbb{R}_{\ge 0})`` with zero slope/curvature at - ``r_inner`` and ``r_outer``, so forces (first derivatives) and the + ``w \in C^3(\mathbb{R}_{\ge 0})`` with zero slope/curvature at both + radii of every pair, so forces (first derivatives) and the force derivatives consumed by second-order training stay continuous across both zone boundaries. - The surrounding infrastructure (``compute_edge_src_gate``) owns the - per-node product reduction and broadcast; this module only encodes - the scalar amplitude shape. + The surrounding infrastructure (``compute_source_gates``) owns the + per-node product reduction and the per-edge leave-one-out product; this + module only encodes the scalar amplitude shape. Parameters ---------- - r_inner : float - Inner radius in Å. At or below this distance ``w = 0``. - r_outer : float - Outer radius in Å. At or above this distance ``w = 1``. + f_inner : float + Inner radius as a fraction of the pair length scale. At or below it + ``w = 0``. + f_outer : float + Outer radius as a fraction of the pair length scale. At or above it + ``w = 1``. Raises ------ ValueError - If ``r_inner <= 0``, ``r_outer <= 0``, or ``r_inner >= r_outer``. + If ``f_inner <= 0``, ``f_outer <= 0``, or ``f_inner >= f_outer``. """ - def __init__(self, r_inner: float, r_outer: float) -> None: - if r_inner <= 0 or r_outer <= 0: - raise ValueError("r_inner and r_outer must be positive") - if r_inner >= r_outer: - raise ValueError(f"r_inner ({r_inner}) must be < r_outer ({r_outer})") - self.r_inner = float(r_inner) - self.r_outer = float(r_outer) + def __init__(self, f_inner: float, f_outer: float) -> None: + if f_inner <= 0 or f_outer <= 0: + raise ValueError("f_inner and f_outer must be positive") + if f_inner >= f_outer: + raise ValueError(f"f_inner ({f_inner}) must be < f_outer ({f_outer})") + self.f_inner = float(f_inner) + self.f_outer = float(f_outer) - def call(self, r: Any) -> Any: + def call(self, r: Any, contact: Any) -> Any: """ Evaluate the C3 switching amplitude. @@ -405,25 +475,114 @@ def call(self, r: Any) -> Any: ---------- r : Array Pair distances with shape (...) or (..., 1) in Å. + contact : Array + Length scale of each pair in Å, broadcastable against ``r``. Returns ------- Array - Switching amplitudes in ``[0, 1]`` with the same shape as input. + Switching amplitudes in ``[0, 1]`` with the same shape as ``r``. """ xp = array_api_compat.array_namespace(r) t = xp.clip( - (r - self.r_inner) / (self.r_outer - self.r_inner), min=0.0, max=1.0 + (r / contact - self.f_inner) / (self.f_outer - self.f_inner), + min=0.0, + max=1.0, ) t2 = t * t t4 = t2 * t2 # h(t) = 35 t^4 - 84 t^5 + 70 t^6 - 20 t^7 (Horner form). # Degree-7 smootherstep: the unique polynomial of this degree that - # hits ``w(r_inner)=0, w(r_outer)=1`` together with C3 flatness at - # both radii. + # hits ``w = 0`` at the inner radius and ``w = 1`` at the outer one + # together with C3 flatness at both. return t4 * (35.0 + t * (-84.0 + t * (70.0 - 20.0 * t))) +class BridgingClamp(NativeOP): + r""" + C4-continuous distance clamp whose slope is the bridging switch. + + The clamp shares the window of :class:`BridgingSwitch` and is its integral: + the distance the descriptor sees moves with the true one at exactly the + rate at which the switch has opened. In the reduced distance ``u = r / s``, + with ``s`` the length scale of the pair:: + + ũ = f_mid if u <= f_inner + ũ = f_mid + (f_outer - f_inner) * S(t) if f_inner < u < f_outer + ũ = u if u >= f_outer + + S(t) = 7 t^5 - 14 t^6 + 10 t^7 - 2.5 t^8, t = (u - f_inner) / (f_outer - f_inner) + + and the clamped distance is ``r̃ = s * ũ``. ``S`` is the antiderivative of + the septic smootherstep ``h`` of the switch, so ``dr̃/dr = h(t)``: zero + below the inner radius, one above the outer radius, never above one in + between. Since ``S(1) = 1/2``, continuity at the outer radius fixes the + frozen value at the window midpoint ``f_mid = (f_inner + f_outer) / 2``. + The vanishing first three derivatives of ``h`` at both ends make ``r̃`` a + C4 function of ``r``. + + Used together with the switch, the clamp keeps the radial features of a + close pair inside the range the training frames cover. The frame filter of + a bridged model keeps pairs down to the window midpoint, the clamped + distance never falls below it, and through the lower half of the window it + rises by less than 7 % of the window width (``S(1/2) = 35/512``), so there + the learned energy follows the switch amplitude. + + Parameters + ---------- + f_inner : float + Inner radius as a fraction of the pair length scale. At or below it the + distance the descriptor sees is frozen at the window midpoint. + f_outer : float + Outer radius as a fraction of the pair length scale. At or above it the + clamp is the identity. + + Raises + ------ + ValueError + If ``f_inner >= f_outer`` or either is non-positive. + """ + + def __init__(self, f_inner: float, f_outer: float) -> None: + if f_inner <= 0 or f_outer <= 0: + raise ValueError("f_inner and f_outer must be positive") + if f_inner >= f_outer: + raise ValueError(f"f_inner ({f_inner}) must be < f_outer ({f_outer})") + self.f_inner = float(f_inner) + self.f_outer = float(f_outer) + + def call(self, r: Any, contact: Any) -> Any: + """ + Apply the distance clamp. + + Parameters + ---------- + r : Array + Pair distances with shape (...) or (..., 1) in Å. + contact : Array + Length scale of each pair in Å, broadcastable against ``r``. + + Returns + ------- + Array + Clamped distances r̃ with the same shape as ``r``. + """ + xp = array_api_compat.array_namespace(r) + width = self.f_outer - self.f_inner + u = r / contact + t = xp.clip((u - self.f_inner) / width, min=0.0, max=1.0) + t2 = t * t + # S(t) = 7 t^5 - 14 t^6 + 10 t^7 - 2.5 t^8 (Horner form), the + # antiderivative of the switch profile with S(0) = 0 and S(1) = 1/2. + rise = t2 * t2 * t * (7.0 + t * (-14.0 + t * (10.0 - 2.5 * t))) + interpolated = contact * ( + window_midpoint(self.f_inner, self.f_outer) + width * rise + ) + # Identity zone: both branches share their first four derivatives at + # the outer radius, so xp.where preserves the continuity there. + return xp.where(u >= self.f_outer, r, interpolated) + + def parse_basis_type(basis_type: str) -> tuple[str, bool]: """ Split a radial basis type into its family and its ``/fix`` flag. diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/so2.py b/deepmd/dpmodel/descriptor/dpa4_nn/so2.py index 323fe3bbfe..2d623fb624 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/so2.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/so2.py @@ -1690,9 +1690,9 @@ def forward_envelope( Reduce the edge messages with the scalar envelope weight. The attention-free path: an envelope-weighted scatter add followed by the - degree normalization. Folding the Source Freeze Propagation Gate into the - envelope keeps the operation count unchanged; ``edge_src_gate`` is ``None`` - outside bridging mode, where the branch disappears. + degree normalization. In bridging mode the envelope already carries the + Source Freeze Propagation Gate, so a muted source drops out of both the + sum and the normalization. Parameters ---------- @@ -1714,12 +1714,8 @@ def forward_envelope( x_message, _ = self.edge_message(x, edge_cache, radial_feat) # (E, D, C_wide) - # === Step 2. Envelope weighting, with the source gate folded in === - edge_weight = edge_cache.edge_env # (E, 1) - edge_src_gate = edge_cache.edge_src_gate - if edge_src_gate is not None: - edge_weight = edge_weight * xp.astype(edge_src_gate, edge_weight.dtype) - x_message = x_message * edge_weight[..., None] # (E, D, C_wide) + # === Step 2. Envelope weighting === + x_message = x_message * edge_cache.edge_env[..., None] # (E, D, C_wide) # === Step 3. Destination reduction and degree normalization === compute_dtype = get_xp_precision(xp, self.compute_precision) @@ -1769,15 +1765,8 @@ def forward_attention( ) # (N, Fa, Ca) # === Step 2. Backend dispatch === - # The fused CUDA operator computes the attention weights itself, so it - # does not serve the bridging mode, whose source gate reshapes the - # softmax normalization. training = getattr(self, "training", False) - run_cuda = ( - self._cuda_conv_fn is not None - and not training - and edge_cache.edge_src_gate is None - ) + run_cuda = self._cuda_conv_fn is not None and not training run_flash = ( self._flash_atten_fn is not None and (not training or self._flash_atten_trains) @@ -2046,12 +2035,12 @@ def attention_weights( """ Build envelope-gated attention weights from the scalar channels. - The softmax takes ``src_weight`` so that the Source Freeze Propagation - Gate enters both the numerator and the denominator. A muted source - (``eta_src = 0``) then drops out of the destination's normalization - entirely, which the frozen-zone invariance requires: post-multiplying the - weights alone would still leak the muted source through the shared - denominator. + The envelope enters both the numerator and the denominator of the + softmax. In bridging mode it carries the Source Freeze Propagation + Gate, so a muted source (``eta_src = 0``) drops out of the + destination's normalization entirely, which the frozen-zone invariance + requires: post-multiplying the weights alone would still leak the muted + source through the shared denominator. Parameters ---------- @@ -2134,7 +2123,6 @@ def _attention_softmax( xp = array_api_compat.array_namespace(attn_logits) device = array_api_compat.device(attn_logits) compute_dtype = get_xp_precision(xp, self.compute_precision) - edge_src_gate = edge_cache.edge_src_gate return segment_envelope_gated_softmax( logits=attn_logits, edge_env=xp.astype(edge_cache.edge_env, compute_dtype), @@ -2144,11 +2132,6 @@ def _attention_softmax( xp, self.adamw_attn_z_bias_raw[...], device=device ), eps=self.eps, - src_weight=( - None - if edge_src_gate is None - else xp.astype(edge_src_gate, compute_dtype) - ), edge_mask=edge_cache.edge_mask, ) diff --git a/deepmd/dpmodel/descriptor/dpa4c.py b/deepmd/dpmodel/descriptor/dpa4c.py index e38d8374da..7dbd9c7f3d 100644 --- a/deepmd/dpmodel/descriptor/dpa4c.py +++ b/deepmd/dpmodel/descriptor/dpa4c.py @@ -58,9 +58,15 @@ from deepmd.dpmodel.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.bridging import ( + check_window_inside_cutoff, +) from deepmd.utils.charge_state import ( validate_charge_state, ) +from deepmd.utils.element_radii import ( + contact_radius_table, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -69,10 +75,13 @@ BaseDescriptor, ) from .dpa4_nn import ( + BridgingClamp, + BridgingSwitch, C3CutoffEnvelope, RadialBasis, SeZMTypeEmbedding, SwiGLUMLP, + pair_contact, resolve_swiglu_hidden_width, ) from .dpa4c_nn import ( @@ -191,6 +200,22 @@ class DescrptDPA4C(NativeOP, BaseDescriptor): spin Reserved for descriptor API compatibility; only ``None`` is supported. Native spin is configured through ``use_spin``. + inner_clamp_f_inner + Inner radius of the zone-bridging window as a fraction of the pair's + length scale, under the name the bridging composition injects into + every DPA4-family descriptor. An edge at or below it leaves the + descriptor entirely, and the pair length the radial features read is + frozen at the window midpoint. Given together with + ``inner_clamp_f_outer`` or not at all. + inner_clamp_f_outer + Outer radius of the zone-bridging window as a fraction of the pair's + length scale. An edge at or beyond it contributes exactly as it does + without a window. + inner_clamp_scale + Length scale the two fractions are measured against: ``"covalent"`` + sizes every pair by the sum of its covalent radii and therefore needs + ``type_map``, ``"absolute"`` gives every pair a unit scale so the + fractions are radii in Å. Raises ------ @@ -198,8 +223,11 @@ class DescrptDPA4C(NativeOP, BaseDescriptor): If ``channels`` or ``lmax`` is not an integer. ValueError If ``rcut``, ``ntypes``, or ``n_radial`` is not positive, if - ``radial_modes`` is negative, or if ``channels`` or ``lmax`` is - outside its supported set. + ``radial_modes`` is negative, if ``channels`` or ``lmax`` is outside + its supported set, if only one bridging radius is given, if + ``inner_clamp_scale`` is neither ``"covalent"`` nor ``"absolute"``, or + if a covalent window is asked for without a ``type_map`` covering every + type. NotImplementedError If ``spin`` is given. """ @@ -250,6 +278,9 @@ def __init__( add_chg_spin_ebd: bool = False, default_chg_spin: list[float] | None = None, spin: None = None, + inner_clamp_f_inner: float | None = None, + inner_clamp_f_outer: float | None = None, + inner_clamp_scale: str = "covalent", ) -> None: # === Step 1. Validate the public architecture contract === if spin is not None: @@ -269,6 +300,16 @@ def __init__( or radial_modes < 0 ): raise ValueError("`radial_modes` must be a non-negative integer.") + if (inner_clamp_f_inner is None) != (inner_clamp_f_outer is None): + raise ValueError( + "`inner_clamp_f_inner` and `inner_clamp_f_outer` must be given " + "together." + ) + if inner_clamp_scale not in ("covalent", "absolute"): + raise ValueError( + "`inner_clamp_scale` must be 'covalent' or 'absolute', got " + f"{inner_clamp_scale!r}." + ) default_chg_spin = ( None if default_chg_spin is None @@ -350,6 +391,35 @@ def __init__( exponent=self._ENVELOPE_EXPONENT, precision=self.precision, ) + # Every edge term carries the envelope and no node feature crosses an + # edge, so one inner switch on the envelope removes a close pair from + # the descriptor as completely as deleting its edge from the graph. + self.bridging_f_inner = ( + None if inner_clamp_f_inner is None else float(inner_clamp_f_inner) + ) + self.bridging_f_outer = ( + None if inner_clamp_f_outer is None else float(inner_clamp_f_outer) + ) + self.bridging_scale = str(inner_clamp_scale) + self.bridging_switch = ( + None + if self.bridging_f_inner is None + else BridgingSwitch(self.bridging_f_inner, self.bridging_f_outer) + ) + self.bridging_clamp = ( + None + if self.bridging_f_inner is None + else BridgingClamp(self.bridging_f_inner, self.bridging_f_outer) + ) + # The table sizes a window, so an unbridged descriptor needs none and + # is not asked for the element symbols the covalent scale reads. + self.contact_radius = None + if self.bridging_f_inner is not None: + table = contact_radius_table( + self.ntypes, self.type_map, self.bridging_scale + ) + check_window_inside_cutoff(table, self.bridging_f_outer, self.rcut) + self.contact_radius = table self.pair_film = OrderedPairFiLM( self.channels, radial_modes=self.radial_modes, @@ -409,6 +479,12 @@ def __init__( self.compress = False self.reinit_exclude(exclude_types) + def pair_contact( + self, center_type: Array, neighbor_type: Array, dtype: Any + ) -> Array: + """Length scale of each edge in Å, shape (E, 1); see :func:`pair_contact`.""" + return pair_contact(self.contact_radius, center_type, neighbor_type, dtype) + # === Descriptor evaluation === def call_graph( @@ -939,7 +1015,19 @@ def build_edge_features( # DPA4C requests the raw radial basis. One explicit C³ envelope gates # the combined radial and type feature, so the scalar edge amplitude # contains exactly one cutoff factor and vanishes smoothly at rcut. - envelope = self.evaluate_cutoff_envelope(distance) * mask + # A bridging window closes the same envelope from the inside: the + # switch reads the true pair length, and the radial features read the + # clamped one, which moves at the rate the switch has opened. The + # direction keeps the true geometry. + inner_switch = None + if self.bridging_switch is not None: + contact = self.pair_contact(center_type, neighbor_type, distance.dtype) + inner_switch = self.bridging_switch.call(distance, contact) + distance = self.bridging_clamp.call(distance, contact) + envelope = self.evaluate_cutoff_envelope(distance) + if inner_switch is not None: + envelope = envelope * inner_switch + envelope = envelope * mask radial_basis = self.radial_basis.call(distance) radial_hidden = self.radial_embedding.call_hidden(radial_basis) radial = self.radial_embedding.call_output(radial_hidden) @@ -1303,9 +1391,10 @@ def structure_signature(self) -> tuple: missing input rather than a property of the shared parameters, so two branches may legitimately default to different charge states. - Branch-local state is deliberately absent. ``exclude_types`` is the - only such field: it configures the pair-exclusion mask, which each - replica keeps for itself. + Branch-local state is deliberately absent: ``exclude_types`` and the + bridging window configure the pair-exclusion mask and the inner + envelope switch, both parameter free, which each replica keeps for + itself. Returns ------- @@ -1736,7 +1825,7 @@ def serialize(self) -> dict: data = { "@class": "Descriptor", "type": "dpa4c", - "@version": 1, + "@version": 2, "rcut": self.rcut, "ntypes": self.ntypes, "channels": self.channels, @@ -1753,6 +1842,9 @@ def serialize(self) -> dict: "add_chg_spin_ebd": self.add_chg_spin_ebd, "default_chg_spin": self.default_chg_spin, "spin": None, + "inner_clamp_f_inner": self.bridging_f_inner, + "inner_clamp_f_outer": self.bridging_f_outer, + "inner_clamp_scale": self.bridging_scale, "spin_channels": (None if self.spin is None else self.spin.serialize()), "charge_spin_embedding": ( None @@ -1798,7 +1890,7 @@ def deserialize(cls, data: dict) -> DescrptDPA4C: Reconstructed descriptor with restored trainable components. """ data = data.copy() - check_version_compatibility(data.pop("@version"), 1, 1) + check_version_compatibility(data.pop("@version"), 2, 1) data.pop("@class") data.pop("type") compression = data.pop("compress", None) diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index 796678f3bd..3619650985 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -80,6 +80,16 @@ def __init__( formatted_descript_list.append(BaseDescriptor(**ii)) else: raise NotImplementedError + # A bridged child fades its fitting output through a readout gate, + # which a concatenated descriptor has no fitting output to apply to. + if any( + getattr(descrpt, "bridging_switch", None) is not None + for descrpt in formatted_descript_list + ): + raise NotImplementedError( + "A hybrid descriptor cannot carry a bridged child; bridge the " + "model with a single descriptor." + ) self.descrpt_list = formatted_descript_list self.numb_descrpt = len(self.descrpt_list) for ii in range(1, self.numb_descrpt): @@ -409,7 +419,7 @@ def call( fparam=fparam, comm_dict=comm_dict, charge_spin=charge_spin, - ) + )[:5] out_descriptor.append(odescriptor) if gr is not None: out_gr.append(gr) diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index fe7b6f3d17..3693cc1c3c 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -123,10 +123,10 @@ class GeneralFitting(NativeOP, BaseFitting): this value will be used as the default value for the frame parameter in the fitting net. """ - # A deployment constant kept out of checkpoints; see + # Deployment constants kept out of checkpoints; see # :meth:`fold_vacuum_reference`. A subclass extends this tuple rather than # replacing it. - CONFIG_DERIVED_ARRAYS = ("vacuum_table",) + CONFIG_DERIVED_ARRAYS = ("vacuum_table", "folded_reference") def __init__( self, @@ -215,6 +215,7 @@ def __init__( else: self.case_embd = None self.vacuum_table = None + self.folded_reference = None if self.default_fparam is not None: if self.numb_fparam > 0: @@ -528,6 +529,7 @@ def change_type_map( self.bias_atom_e = self.bias_atom_e[remap_index] # the stored references belong to the old type map self.vacuum_table = None + self.folded_reference = None def __setitem__(self, key: str, value: Any) -> None: if key in ["bias_atom_e"]: @@ -892,12 +894,14 @@ def fold_vacuum_reference(self, vacuum_descriptor: Array) -> None: Under uniform conditioning the reference output of an atom is a constant of its type, so subtracting it from ``bias_atom_e`` yields the same outputs as the referenced call and the option is switched - off. With frame or atomic parameters the reference output varies - between atoms, so the vacuum descriptor is stored instead and the call - evaluates the references from the stored table. The table is a - deployment constant: an exported model bakes it, checkpoints leave it - out, and a fitting loaded from a checkpoint takes the reference from - the descriptor again. + off; the subtracted outputs are kept as ``folded_reference`` so that + :meth:`readout_reference` still knows the bias the fold moved. With + frame or atomic parameters the reference output varies between atoms, + so the vacuum descriptor is stored instead and the call evaluates the + references from the stored table. Both are deployment constants: an + exported model bakes them, checkpoints leave them out, and a fitting + loaded from a checkpoint takes the reference from the descriptor + again. Parameters ---------- @@ -910,12 +914,36 @@ def fold_vacuum_reference(self, vacuum_descriptor: Array) -> None: if not self.uniform_conditioning(): self.vacuum_table = to_numpy_array(vacuum_descriptor) return - reference = to_numpy_array(self.vacuum_property(vacuum_descriptor)) - self["bias_atom_e"] = to_numpy_array(self.bias_atom_e) - reference.astype( + reference = to_numpy_array(self.vacuum_property(vacuum_descriptor)).astype( GLOBAL_NP_FLOAT_PRECISION ) + self.folded_reference = reference + self["bias_atom_e"] = to_numpy_array(self.bias_atom_e) - reference self.vacuum_ref = False + def readout_reference(self) -> Array: + """Per-type constant the readout gate of a bridging window fades to. + + It is the fitting bias ``bias_atom_e`` of every type, which under the + vacuum reference is the output of an isolated atom of that type. A + fold under uniform conditioning moves the reference outputs out of the + bias; the stored copy restores the bias the gate reads, so the gate + is the same function of the atom before and after the fold. + + Returns + ------- + Array + The reference of every type with shape (ntypes, dim_out). + """ + if self.folded_reference is None: + return self.bias_atom_e + xp = array_api_compat.array_namespace(self.bias_atom_e) + return self.bias_atom_e + xp.asarray( + self.folded_reference, + dtype=self.bias_atom_e.dtype, + device=array_api_compat.device(self.bias_atom_e), + ) + def _call_common( self, descriptor: Array, @@ -926,6 +954,7 @@ def _call_common( fparam: Array | None = None, aparam: Array | None = None, vacuum_descriptor: Array | None = None, + node_gate: Array | None = None, ) -> dict[str, Array]: """Calculate the fitting. @@ -951,6 +980,9 @@ def _call_common( vacuum_descriptor The descriptor of an isolated atom of every type, required by ``vacuum_ref``. shape: ntypes x nd + node_gate + Per-atom source gate of a bridged descriptor; the learned part of + the output fades with it. shape: nf x nloc x 1 """ xp = array_api_compat.array_namespace(descriptor, atype) @@ -1050,6 +1082,27 @@ def _call_common( ), (nf, nloc, net_dim_out), ) + # === Step 3. Readout gate of a bridging window === + # The learned part of the output is its deviation from the bias of + # the type (see :meth:`readout_reference`), which is the output of an + # isolated atom whenever the fitting references its atoms. A bridged + # descriptor hands over the per-atom source gate, and the learned part + # fades with it: an atom of a frozen pair keeps only the bias of its + # type. + if node_gate is not None: + reference = xp.reshape( + xp.take( + xp.asarray( + self.readout_reference(), + dtype=outs.dtype, + device=array_api_compat.device(outs), + ), + xp.reshape(atype, (-1,)), + axis=0, + ), + (nf, nloc, net_dim_out), + ) + outs = reference + xp.astype(node_gate, outs.dtype) * (outs - reference) # nf x nloc exclude_mask = self.emask.build_type_exclude_mask(atype) exclude_mask = xp.astype(exclude_mask, xp.bool) @@ -1070,6 +1123,7 @@ def call_graph( fparam: Array | None = None, aparam: Array | None = None, vacuum_descriptor: Array | None = None, + node_gate: Array | None = None, ) -> dict[str, Array]: """Graph-native (flat node axis) fitting forward. @@ -1098,6 +1152,9 @@ def call_graph( vacuum_descriptor the descriptor of an isolated atom of every type, required by ``vacuum_ref``. ntypes x nd + node_gate + per-atom source gate of a bridged descriptor; the learned part of + the output fades with it. N Returns ------- @@ -1130,6 +1187,9 @@ def call_graph( if vacuum_descriptor is None else {"vacuum_descriptor": vacuum_descriptor} ) + gate_kwargs = ( + {} if node_gate is None else {"node_gate": xp.reshape(node_gate, (n, 1, 1))} + ) ret = self.__call__( d1, a1, @@ -1139,5 +1199,6 @@ def call_graph( fparam=fparam, aparam=ap1, **vacuum_kwargs, + **gate_kwargs, ) return {kk: xp.reshape(vv, (n, *vv.shape[2:])) for kk, vv in ret.items()} diff --git a/deepmd/dpmodel/fitting/invar_fitting.py b/deepmd/dpmodel/fitting/invar_fitting.py index 57c9eec185..9268ebfe85 100644 --- a/deepmd/dpmodel/fitting/invar_fitting.py +++ b/deepmd/dpmodel/fitting/invar_fitting.py @@ -241,6 +241,7 @@ def call( fparam: Array | None = None, aparam: Array | None = None, vacuum_descriptor: Array | None = None, + node_gate: Array | None = None, ) -> dict[str, Array]: """Calculate the fitting. @@ -266,6 +267,9 @@ def call( vacuum_descriptor The descriptor of an isolated atom of every type, required by ``vacuum_ref``. shape: ntypes x nd + node_gate + Per-atom source gate of a bridged descriptor; the learned part of + the output fades with it. shape: nf x nloc x 1 """ return self._call_common( @@ -277,4 +281,5 @@ def call( fparam, aparam, vacuum_descriptor=vacuum_descriptor, + node_gate=node_gate, ) diff --git a/deepmd/dpmodel/model/dp_linear_model.py b/deepmd/dpmodel/model/dp_linear_model.py index 2255edbede..662c5f431c 100644 --- a/deepmd/dpmodel/model/dp_linear_model.py +++ b/deepmd/dpmodel/model/dp_linear_model.py @@ -3,6 +3,9 @@ composition (twin of ``deepmd.pt_expt.model.dp_linear_model``). """ +from collections.abc import ( + Callable, +) from typing import ( Any, ) @@ -49,6 +52,12 @@ def __init__( DPModelCommon.__init__(self) DPLinearModel_.__init__(self, *args, **kwargs) + def compute_fitting_input_stats( + self, sampled_func: Callable[[], list[dict]] + ) -> None: + """Recompute the input statistics of every learned sub-model's fitting net.""" + self.atomic_model.compute_fitting_input_stats(sampled_func) + @classmethod def update_sel( cls, diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py index 5a1a253596..7f3f8b9705 100644 --- a/deepmd/dpmodel/model/dp_model.py +++ b/deepmd/dpmodel/model/dp_model.py @@ -1,5 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from deepmd.dpmodel.descriptor.base_descriptor import ( BaseDescriptor, @@ -58,3 +61,15 @@ def get_fitting_net(self) -> BaseFitting: def get_descriptor(self) -> BaseDescriptor: """Get the descriptor.""" return self.atomic_model.descriptor + + def compute_fitting_input_stats( + self, sampled_func: Callable[[], list[dict]] + ) -> None: + """Recompute the input statistics of the fitting network from a sample. + + Parameters + ---------- + sampled_func + The lazy sampled function to get data frames from different data systems. + """ + self.get_fitting_net().compute_input_stats(sampled_func) diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 3762f9304f..66d32f8f8b 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -87,9 +87,9 @@ def get_linear_model(data: dict) -> BaseModel: models; ``pairtab`` children build as pair-tabulation atomic models; an ``inner_potential`` child builds the analytical bridging term. The composition is the ONE owner of the bridging coupling: it derives the - learned sibling descriptor's ``inner_clamp_r_inner``/``_outer`` from - the ``inner_potential`` child's ``r_inner``/``r_outer``, so the radii - are written once in the config (issue #5948, task 2). + learned sibling descriptor's window from the ``inner_potential`` + child's, so the window is written once in the config (issue #5948, + task 2). A top-level ``spin`` section (scheme ``native``) wraps the composed atomic model as a :class:`NativeSpinEnergyModel`, with ``use_spin`` diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py index d49ad1d9df..38a4f4b9e4 100644 --- a/deepmd/dpmodel/model/model_factory.py +++ b/deepmd/dpmodel/model/model_factory.py @@ -11,10 +11,12 @@ ) from deepmd.utils.bridging import ( + resolve_bridging_window, route_canonical_learned_options, ) from deepmd.utils.spin import ( Spin, + env_protection_descriptors, ) ModelBuilder = Callable[[dict], Any] @@ -147,9 +149,9 @@ def get_linear_atomic_model( atomic models; an ``inner_potential`` child builds the analytical bridging term. The composition is the ONE owner of the bridging coupling: it derives the learned sibling descriptor's - ``inner_clamp_r_inner``/``_outer`` from the ``inner_potential`` - child's ``r_inner``/``r_outer``, so the radii are written once in the - config (issue #5948). + ``inner_clamp_f_inner``/``_f_outer``/``_scale`` from the + ``inner_potential`` child's window, so the window is written once in + the config (issue #5948). Parameters ---------- @@ -266,21 +268,27 @@ def get_linear_atomic_model( learned_descriptor_type = str( children[learned_indices[0]]["descriptor"].get("type", "dpa4") ) - if learned_descriptor_type not in ("dpa4", "DPA4", "sezm", "SeZM"): - # same family restriction as the pt builder: the clamp window - # below only exists on DPA4/SeZM descriptors, so any other + if learned_descriptor_type not in ( + "dpa4", + "DPA4", + "sezm", + "SeZM", + "dpa4c", + "DPA4C", + ): + # same family restriction as the pt builder: the bridging window + # below only exists on DPA4-family descriptors, so any other # family would die on an obscure unknown-kwarg TypeError raise NotImplementedError( f"The {backend_name} backend implements `inner_potential` " - "bridging only for the DPA4/SeZM descriptor family, but got " - f"{learned_descriptor_type!r}." + "bridging only for the DPA4 descriptor family (DPA4/SeZM and " + f"DPA4C), but got {learned_descriptor_type!r}." ) # The composition derives the sibling descriptor's clamp window from - # the inner_potential child: one source of truth for the radii. + # the inner_potential child: one source of truth for the window. inner_cfg = children[inner_indices[0]] learned_descriptor = children[learned_indices[0]]["descriptor"] - learned_descriptor["inner_clamp_r_inner"] = float(inner_cfg.get("r_inner", 0.5)) - learned_descriptor["inner_clamp_r_outer"] = float(inner_cfg.get("r_outer", 0.8)) + learned_descriptor.update(resolve_bridging_window(inner_cfg)) route_canonical_learned_options(data, children[learned_indices[0]]) preset_out_bias = children[learned_indices[0]].get("preset_out_bias") @@ -309,7 +317,15 @@ def get_linear_atomic_model( fitting_base=fitting_base, backend_name=backend_name, ) - child = atomic_model(descriptor, fitting, type_map=sub["type_map"]) + # The composition computes the output bias, so its preset + # decides whether the fitting of a child references the + # isolated atoms. + child = atomic_model( + descriptor, + fitting, + type_map=sub["type_map"], + preset_out_bias=preset_out_bias, + ) built[i] = child else: if sub.get("type") != "pairtab": @@ -366,7 +382,8 @@ def get_spin_model( data["atom_exclude_types"] = spin.get_atom_exclude_types( exclude_types=data.get("atom_exclude_types") ) - data["descriptor"].setdefault("env_protection", 1e-6) + for descriptor in env_protection_descriptors(data["descriptor"]): + descriptor.setdefault("env_protection", 1e-6) if data["descriptor"]["type"] == "se_e2_a": data["descriptor"]["sel"] += data["descriptor"]["sel"] backbone_model = standard_model_factory(data) diff --git a/deepmd/dpmodel/train/trainer.py b/deepmd/dpmodel/train/trainer.py index 531bc7d07c..22cb543ced 100644 --- a/deepmd/dpmodel/train/trainer.py +++ b/deepmd/dpmodel/train/trainer.py @@ -84,7 +84,7 @@ def change_model_out_bias( new_bias = deepcopy(model.get_out_bias()) if recompute_input_stats and bias_adjust_mode == "set-by-statistic": - model.get_fitting_net().compute_input_stats(sample_func) + model.compute_fitting_input_stats(sample_func) model_type_map = model.get_type_map() log.info( diff --git a/deepmd/dpmodel/utils/dist_check.py b/deepmd/dpmodel/utils/dist_check.py index 3c5eb25681..e8218d8186 100644 --- a/deepmd/dpmodel/utils/dist_check.py +++ b/deepmd/dpmodel/utils/dist_check.py @@ -1,48 +1,188 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Minimum pairwise distance check for frame validity filtering.""" +r"""Pair-clearance check for frame validity filtering. + +A frame is admissible when every atom pair keeps the clearance the +training-frame filter prescribes. That clearance is a fraction of each pair's +own length scale, so the check is scale-free: the quantity compared against +one is the margin + +.. math:: + + m = \min_{i np.ndarray: + """ + Tabulate the half-threshold each type contributes to a pair. + + A pair's threshold is ``half[a] + half[b]``. Under the ``"covalent"`` scale + that sum is the given fraction of the pair's covalent bond length, so one + fraction describes every element combination; under ``"absolute"`` every + type carries the unit radius, the sum is the fraction itself, and the + margin reduces to the pair distance divided by a threshold in Å. + + Parameters + ---------- + fraction : float + Filter radius as a fraction of the pair's length scale, or directly in + Å under the ``"absolute"`` scale. + scale : str + Either ``"covalent"`` or ``"absolute"``. + type_map : list[str], optional + Element symbols, required by the ``"covalent"`` scale. + ntypes : int, optional + Number of types; taken from ``type_map`` when omitted. + + Returns + ------- + np.ndarray + Half-thresholds in Å with shape ``(ntypes,)``. + + Raises + ------ + ValueError + If the fraction is not positive, if the scale is unknown, if the + covalent scale is requested without a type map, or if the map does not + cover every type. + """ + if fraction <= 0.0: + raise ValueError(f"A pair window needs a positive radius, got {fraction}.") + if scale not in ("covalent", "absolute"): + raise ValueError( + f"Unknown pair length scale {scale!r}; expected 'covalent' or 'absolute'." + ) + if scale == "absolute": + if ntypes is None: + if type_map is None: + raise ValueError( + "An absolute pair length scale needs either `ntypes` or a " + "`type_map` to size its table." + ) + ntypes = len(type_map) + return np.full(ntypes, fraction * UNIT_CONTACT_RADIUS, dtype=np.float64) + if type_map is None or (ntypes is not None and len(type_map) != ntypes): + raise ValueError( + "A covalent pair length scale needs a `type_map` covering all " + f"{ntypes if ntypes is not None else 'model'} types." + ) + return fraction * covalent_radii_from_type_map(type_map) + + +@functools.lru_cache(maxsize=8) +def requirement_half_thresholds( + fraction: float, scale: str, type_map: tuple[str, ...] | None, ntypes: int +) -> np.ndarray: + """ + Resolve and retain the half-thresholds a derived margin measures against. + + A reader scans every frame and every batch against one window and one type + map, so the table is resolved once per combination instead of once per + scan. The result is shared and therefore read-only. + + Parameters + ---------- + fraction : float + Filter radius, as a fraction of the pair length scale. + scale : str + Either ``"covalent"`` or ``"absolute"``. + type_map : tuple[str, ...], optional + Type names, required by the ``"covalent"`` scale. + ntypes : int + Number of model types. + + Returns + ------- + np.ndarray + Half-thresholds in Å with shape ``(ntypes,)``. + """ + table = pair_half_thresholds( + fraction, + scale, + type_map=None if type_map is None else list(type_map), + ntypes=ntypes, + ) + table.flags.writeable = False + return table -def compute_min_pair_dist_single( +def compute_min_pair_margin_single( coord: np.ndarray, box: np.ndarray | None, atype: np.ndarray, - stop_below: float | None = None, + half: np.ndarray, + screened: bool = False, ) -> float: - """Compute the minimum pairwise atomic distance for a single frame. + """ + Compute the minimum pair margin of a single frame. + + Pairs are visited in bounded row blocks, every one of them unless the scan + is screened. Parameters ---------- coord : np.ndarray Atomic coordinates, flattened with shape (natoms * 3,) - or reshaped as (natoms, 3). + or reshaped as (natoms, 3), in Å. box : np.ndarray or None Box vectors with shape (9,) for PBC, or None for non-PBC. atype : np.ndarray Atom types with shape (natoms,). Virtual atoms (type < 0) - are excluded from the distance check. - stop_below : float or None - Optional early-stop threshold. If a block has any pair closer - than this value, the block minimum is returned immediately. + are excluded from the check. + half : np.ndarray + Half-threshold of every type with shape (ntypes,) in Å. + screened : bool + Whether the scan may stop at a margin of one. A caller that only asks + whether the frame clears its window learns that from the first pair + inside it, and the remaining pairs of a failing frame cost the scan its + whole cost on the largest frames. Returns ------- float - Minimum pairwise distance. Returns inf if fewer than 2 - real atoms exist. + Minimum pair margin, or inf if fewer than 2 real atoms exist. A + screened scan returns the exact margin for a frame that clears its + window and some margin below one for a frame that does not. """ coord = coord.reshape(-1, 3) # === Step 1. Filter out virtual atoms === real_mask = atype.ravel() >= 0 real_coord = coord[real_mask] + real_half = np.asarray(half, dtype=np.float64)[atype.ravel()[real_mask]] n_real = real_coord.shape[0] if n_real < 2: return float("inf") @@ -55,14 +195,10 @@ def compute_min_pair_dist_single( cell = None inv_cell = None - # === Step 3. Compute distances in bounded row blocks === + # === Step 3. Compute margins in bounded row blocks === block_size = max(1, min(n_real, _MIN_PAIR_DIST_BLOCK_PAIRS // n_real)) - min_dist_sq = float("inf") - stop_dist_sq = ( - float(stop_below) * float(stop_below) - if stop_below is not None and stop_below > 0.0 - else None - ) + min_margin_sq = float("inf") + stop_below_sq = 1.0 if screened else 0.0 for start in range(0, n_real, block_size): stop = min(start + block_size, n_real) diff = real_coord[np.newaxis, :, :] - real_coord[start:stop, np.newaxis, :] @@ -72,13 +208,399 @@ def compute_min_pair_dist_single( frac_diff -= np.round(frac_diff) diff = frac_diff @ cell - dist_sq = np.sum(diff * diff, axis=-1) + threshold = real_half[np.newaxis, :] + real_half[start:stop, np.newaxis] + margin_sq = np.sum(diff * diff, axis=-1) / (threshold * threshold) rows = np.arange(stop - start, dtype=np.int64) - dist_sq[rows, start + rows] = np.inf - min_dist_sq = min(min_dist_sq, float(dist_sq.min())) - if min_dist_sq == 0.0 or ( - stop_dist_sq is not None and min_dist_sq < stop_dist_sq - ): + margin_sq[rows, start + rows] = np.inf + min_margin_sq = min(min_margin_sq, float(margin_sq.min())) + # A screened scan is settled by the first pair strictly inside its + # window; an exhaustive one can stop only at zero, the margin's floor. + if min_margin_sq < stop_below_sq or min_margin_sq == 0.0: break - return float(np.sqrt(min_dist_sq)) + return float(np.sqrt(min_margin_sq)) + + +def compute_min_pair_margin_batch( + coord: np.ndarray, + box: np.ndarray | None, + atype: np.ndarray, + half: np.ndarray, + n_node: np.ndarray | None = None, + screened: bool = False, +) -> np.ndarray: + """ + Compute the minimum pair margin of every frame of a batch. + + A batch is scanned in one pass over the frames that share an atom count + rather than one pass per frame, and the scan of a frame visits the atom + pairs that lie close along one lattice direction instead of every pair + (see :func:`_min_pair_margin_block`). + + Parameters + ---------- + coord : np.ndarray + Atomic coordinates, either `(nframes, natoms, 3)` (equivalently + `(nframes, natoms * 3)`) for a rectangular batch, or `(nnodes, 3)` + (equivalently `(nnodes * 3,)`) on the flat node axis of a ragged one, + in Å. + box : np.ndarray or None + Box vectors with shape `(nframes, 9)` for periodic frames, or None. + atype : np.ndarray + Atom types, `(nframes, natoms)` or `(nnodes,)`. Virtual atoms + (type < 0) are excluded from the check. + half : np.ndarray + Half-threshold of every type with shape `(ntypes,)` in Å. + n_node : np.ndarray or None + Per-frame atom counts of a ragged batch, with shape `(nframes,)`. + None marks a rectangular batch. + screened : bool + Bound the scan of a frame at a margin of one, which is all the frame + filter reads. A frame holding a pair inside its window then keeps its + exact margin, and a frame clearing it carries a margin above one that + it reaches. + + Returns + ------- + np.ndarray + Margin of every frame, with shape `(nframes,)`. A frame without a pair + of real atoms carries inf. + """ + half = np.asarray(half, dtype=np.float64) + boxes = None if box is None else np.asarray(box).reshape(-1, 9) + if n_node is None: + atype = np.asarray(atype).reshape(len(atype), -1) + nframes, natoms = atype.shape + coord = np.asarray(coord).reshape(nframes, natoms, 3) + real = atype >= 0 + if natoms > 0 and real.all(): + return _scan_uniform(coord, half[atype], boxes, screened) + counts = real.sum(axis=1) + positions = coord[real] + halves = half[atype[real]] + else: + n_node = np.asarray(n_node).reshape(-1) + coord = np.asarray(coord).reshape(-1, 3) + atype = np.asarray(atype).reshape(-1) + real = atype >= 0 + counts = np.bincount( + np.repeat(np.arange(n_node.shape[0], dtype=np.int64), n_node)[real], + minlength=n_node.shape[0], + ) + positions = coord[real] + halves = half[atype[real]] + return _scan_by_atom_count(positions, halves, boxes, counts, screened) + + +def _scan_by_atom_count( + positions: np.ndarray, + halves: np.ndarray, + boxes: np.ndarray | None, + counts: np.ndarray, + screened: bool, +) -> np.ndarray: + """Scan frames of real atoms laid out on one flat node axis. + + Frames of equal atom count form the rectangular blocks the scan works on. + A batch of the LMDB reader normally holds one count, and a padded + rectangular batch reduces to one as soon as its filler atoms are left out, + so this usually yields a single block; a block too small to pay for the + batched scan is served frame by frame instead. + + Parameters + ---------- + positions : np.ndarray + Coordinates of the real atoms of every frame in order, `(nnodes, 3)`. + halves : np.ndarray + Half-threshold of each of those atoms, `(nnodes,)` in Å. + boxes : np.ndarray or None + Box vectors with shape `(nframes, 9)`, or None for open boundaries. + counts : np.ndarray + Real atom count of every frame, with shape `(nframes,)`. + screened : bool + Whether the scan is bounded at a margin of one. + + Returns + ------- + np.ndarray + Margin of every frame, with shape `(nframes,)`. + """ + offsets = np.concatenate([[0], np.cumsum(counts)[:-1]]) + margins = np.full(counts.shape[0], np.inf, dtype=np.float64) + for size in np.unique(counts): + if size < 2: + continue + rows = np.flatnonzero(counts == size) + block_boxes = None if boxes is None else boxes[rows] + if rows.shape[0] < _MIN_PAIR_DIST_BATCH_FRAMES: + # The half-thresholds of these atoms are already resolved, so each + # atom stands for its own type and the lookup becomes the identity. + all_real = np.arange(size, dtype=np.int64) + for index, row in enumerate(rows): + start = offsets[row] + cell = None if block_boxes is None else block_boxes[index] + margins[row] = compute_min_pair_margin_single( + positions[start : start + size], + None if cell is None or not cell.any() else cell, + all_real, + halves[start : start + size], + screened=screened, + ) + continue + nodes = offsets[rows][:, np.newaxis] + np.arange(size, dtype=np.int64) + margins[rows] = _scan_uniform( + positions[nodes], halves[nodes], block_boxes, screened + ) + return margins + + +def _scan_uniform( + coord: np.ndarray, + halves: np.ndarray, + boxes: np.ndarray | None, + screened: bool, +) -> np.ndarray: + """Scan frames of one atom count, all of whose atoms are real. + + A cell of all zeros marks a frame without periodicity, which the readers + may mix with periodic ones; the two kinds are scanned apart because the + minimum image of a frame only exists with a cell. + + Parameters + ---------- + coord : np.ndarray + Coordinates with shape `(nframes, natoms, 3)`. + halves : np.ndarray + Half-threshold of every atom, `(nframes, natoms)` in Å. + boxes : np.ndarray or None + Box vectors with shape `(nframes, 9)`, or None for open boundaries. + screened : bool + Whether the scan is bounded at a margin of one. + + Returns + ------- + np.ndarray + Margin of every frame, with shape `(nframes,)`. + """ + nframes, natoms = coord.shape[:2] + if natoms < 2: + return np.full(nframes, np.inf, dtype=np.float64) + if boxes is None: + return _min_pair_margin_block(coord, halves, None, screened) + periodic = ~np.all(boxes == 0.0, axis=1) + if periodic.all(): + return _min_pair_margin_block(coord, halves, boxes, screened) + margins = np.empty(nframes, dtype=np.float64) + if periodic.any(): + margins[periodic] = _min_pair_margin_block( + coord[periodic], halves[periodic], boxes[periodic], screened + ) + margins[~periodic] = _min_pair_margin_block( + coord[~periodic], halves[~periodic], None, screened + ) + return margins + + +def _min_pair_margin_block( + coord: np.ndarray, + halves: np.ndarray, + box: np.ndarray | None, + screened: bool = False, +) -> np.ndarray: + """Minimum pair margin of frames of equal atom count. + + The atoms of every frame are ordered along one direction -- the lattice + planes that lie furthest apart, or the widest Cartesian spread without a + cell -- and each atom is paired with its successors in that order. A pair + can only improve a frame's margin `m` when its distance falls below + `m` times its own threshold, hence below `m` times twice the largest + half-threshold the frame carries; the separation along the ordering + direction only grows with the distance in the order, so a frame is + finished as soon as every atom's window reaches that far. The scan + therefore visits the successors that lie within one slab instead of every + pair, and its bound is a bound on the margin rather than on the distance: + the pair that minimizes the margin need not be the pair that minimizes the + distance, and the scan never assumes it is. + + Under periodicity the order is cyclic and every unordered pair is met from + both of its atoms, at the windows `shift` and `natoms - shift`. That is + what lets the advance be measured in the direction of the order: the pair + whose periodic image is close the other way around is met from its other + atom, where the advance is short. + + Parameters + ---------- + coord : np.ndarray + Coordinates with shape `(nframes, natoms, 3)`. Every atom counts. + halves : np.ndarray + Half-threshold of every atom, `(nframes, natoms)` in Å. + box : np.ndarray or None + Box vectors with shape `(nframes, 9)`, or None for open boundaries. + screened : bool + Whether the scan is bounded at a margin of one. + + Returns + ------- + np.ndarray + Margin of every frame, with shape `(nframes,)`. + """ + nframes, natoms = coord.shape[:2] + if natoms < 2: + return np.full(nframes, np.inf, dtype=np.float64) + + coord = np.ascontiguousarray(coord, dtype=np.float64) + if box is None: + # Open boundaries: the order follows a Cartesian axis and a separation + # along it is already a distance. + position = coord + metric = None + spacing = np.ones((nframes, 3), dtype=np.float64) + extent = np.ptp(coord, axis=1) + else: + cell = np.asarray(box, dtype=np.float64).reshape(nframes, 3, 3) + inverse = np.linalg.inv(cell) + position = coord @ inverse + position -= np.floor(position) + metric = cell @ np.swapaxes(cell, -1, -2) + # Perpendicular width of each family of lattice planes: a fractional + # separation of one along an axis is this distance in space. + spacing = 1.0 / np.linalg.norm(np.swapaxes(inverse, -1, -2), axis=-1) + extent = spacing + + # One direction serves the whole block, so it is the one that separates + # the atoms of its worst frame the most. + axis = int(np.argmax(extent.min(axis=0))) + order = np.argsort(position[:, :, axis], axis=1) + rows = np.arange(nframes, dtype=np.int64)[:, np.newaxis] + position = position[rows, order] + # Half-thresholds follow the atoms into the order, and the largest of a + # frame turns its margin bound into a distance the scan can measure. + halves = np.asarray(halves, dtype=np.float64)[rows, order] + widest = (2.0 * halves.max(axis=1)) ** 2 + + minima = np.full(nframes, np.inf, dtype=np.float64) + active = np.arange(nframes, dtype=np.int64) + shift = 1 + cyclic = box is not None + while active.size and shift < natoms: + window = position[active] + threshold = halves[active] + partner = np.roll(window, -shift, axis=1) if cyclic else window[:, shift:] + head = window if cyclic else window[:, : natoms - shift] + partner_half = ( + np.roll(threshold, -shift, axis=1) if cyclic else threshold[:, shift:] + ) + head_half = threshold if cyclic else threshold[:, : natoms - shift] + separation = partner - head + if cyclic: + separation -= np.round(separation) + # The quadratic form is taken in two steps so that the metric + # enters as a matrix product, which reaches BLAS; the three-operand + # contraction falls back to an unblocked loop and costs twice as + # much for the same result. + squared = np.einsum("fnx,fnx->fn", separation @ metric[active], separation) + else: + squared = np.sum(separation * separation, axis=-1) + pair_threshold = head_half + partner_half + minima[active] = np.minimum( + minima[active], (squared / (pair_threshold * pair_threshold)).min(axis=1) + ) + + pending = minima[active] + if screened: + # No frame needs its window to travel further than one margin's + # worth: beyond that only a pair already outside its own window + # could still be found, which cannot change the frame's side. A + # frame that has reached one therefore keeps the exact margin, + # because from then on this bound is its running minimum. + pending = np.minimum(pending, 1.0) + advance = partner[:, :, axis] - head[:, :, axis] + if cyclic: + advance -= np.floor(advance) + reach = advance * spacing[active, axis, np.newaxis] + settled = (reach**2 >= (pending * widest[active])[:, np.newaxis]).all(axis=1) + active = active[~settled] + shift += 1 + return np.sqrt(minima) + + +def pair_margin_frame_mask( + batch: dict[str, Any], + enabled: bool, +) -> np.ndarray | None: + """Mark the frames of a batch whose atom pairs all clear their window. + + The filter radius the requirement was registered with is already spent: + it sized the per-type half-thresholds the margin was measured against. + What is left here is the same comparison against one for every length + scale, so the only decision this function still takes from the + configuration is whether the filter runs at all. + + Parameters + ---------- + batch : dict[str, Any] + Normalized NumPy batch carrying the derived `pair_margin` field, + which holds each frame's margin against its own pair thresholds. + enabled : bool + Whether the pair-clearance filter is configured for this data. + + Returns + ------- + np.ndarray or None + Boolean mask with shape (nframes,), or None when the filter is + disabled. + + Raises + ------ + RuntimeError + If the batch carries a placeholder instead of a derived margin, + which would otherwise accept every frame silently. + """ + if not enabled: + return None + if not batch.get("find_pair_margin", False): + raise RuntimeError( + "The pair-clearance filter is enabled, but the data source did not " + "derive the pair margin of the batch." + ) + return batch["pair_margin"].reshape(-1) >= 1.0 + + +def select_frames(batch: dict[str, Any], frame_mask: np.ndarray) -> dict[str, Any]: + """Keep the masked frames of a normalized NumPy batch. + + A rectangular batch leads every per-frame and per-atom array with the + frame axis. A ragged batch, which carries `n_node`, leads its per-frame + arrays with the frame axis as well but concatenates its per-atom arrays + frame by frame on one flat node axis, with a whole number of rows per atom; + those are sliced by the node mask that the frame mask expands to. + + Parameters + ---------- + batch : dict[str, Any] + Normalized NumPy batch, rectangular or ragged. + frame_mask : np.ndarray + Boolean mask with shape (nframes,). + + Returns + ------- + dict[str, Any] + New batch in which every array on the frame axis or on the flat node + axis is sliced; all other entries pass through. + """ + nframes = frame_mask.shape[0] + n_node = batch.get("n_node") + nnodes = 0 if n_node is None else int(n_node.sum()) + node_mask = None if n_node is None else np.repeat(frame_mask, n_node) + + def select(value: Any) -> Any: + if not isinstance(value, np.ndarray) or value.ndim == 0: + return value + # Frames of one atom make the two axes coincide, and the two masks too. + if value.shape[0] == nframes: + return value[frame_mask] + if nnodes > 0 and value.shape[0] % nnodes == 0: + return value[np.repeat(node_mask, value.shape[0] // nnodes)] + return value + + return {key: select(value) for key, value in batch.items()} diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index 897292eeb7..0d7ea42a89 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -43,7 +43,9 @@ import numpy as np from deepmd.dpmodel.utils.dist_check import ( - compute_min_pair_dist_single, + compute_min_pair_margin_batch, + compute_min_pair_margin_single, + requirement_half_thresholds, ) from deepmd.env import ( GLOBAL_ENER_FLOAT_PRECISION, @@ -125,7 +127,7 @@ "charge_spin", "natoms", "real_natoms_vec", - "min_pair_dist", + "pair_margin", } ) @@ -688,6 +690,9 @@ class LmdbDecodeConfig: Registered data requirements keyed by field name. dataset Dataset identifier used in frame-level diagnostics. + type_map + Element symbols of the model types, from which a derived field sizes + the per-type length scale it measures against. frame_transform Optional callable applied to every decoded frame, as ``transform(frame, frame_index)``. Self-supervised objectives corrupt @@ -701,6 +706,7 @@ class LmdbDecodeConfig: type_remap: np.ndarray | None data_requirements: dict[str, Any] dataset: str = "" + type_map: list[str] | None = None frame_transform: Callable[[dict[str, Any], int], dict[str, Any]] | None = None @@ -845,6 +851,92 @@ def _compute_frame_natoms(atype: np.ndarray, ntypes: int) -> np.ndarray: return natoms +def _resolve_derived_batch( + batch: dict[str, Any], + key: str, + requirement: Any, + config: LmdbDecodeConfig, + layout: "BatchLayout", +) -> None: + """Resolve one derived field from the assembled arrays of a batch. + + The batch holds the same structural data as its frames, and taking the + field here lets the geometric scan cross the frames of the batch in one + pass. The value settles the frame against the window the requirement + declares, exactly as the per-frame + :func:`_resolve_derived_requirement` does, and equals the minimum pair + margin whenever that margin decides the frame; a frame the window accepts + may carry a larger margin that it reaches, because the scan stops once + the side is known. + + Parameters + ---------- + batch : dict[str, Any] + Assembled batch, modified in place. + key : str + Name of the derived field. + requirement : Any + Its data requirement. + config : LmdbDecodeConfig + Decoder state, which fixes the field dtype and the pair length scale. + layout : BatchLayout + Where each frame's per-atom rows are, which tells a ragged batch from + a rectangular one. + """ + if key != "pair_margin": + raise ValueError(f"Unsupported derived LMDB field {key!r}") + + coord = batch.get("coord") + atype = batch.get("atype") + if not isinstance(coord, np.ndarray) or not isinstance(atype, np.ndarray): + batch.pop(key, None) + batch[f"find_{key}"] = np.float32(0.0) + return + + # A frame without periodicity carries a cell of zeros, which the scan + # reads frame by frame; the batch may hold both kinds. + margins = compute_min_pair_margin_batch( + coord, + batch.get("box"), + atype, + _derived_half_thresholds(requirement, config), + n_node=layout.n_node if layout.ragged else None, + screened=True, + ) + batch[key] = margins.reshape(-1, 1).astype( + _resolve_frame_dtype(config, key), copy=False + ) + batch[f"find_{key}"] = np.float32(1.0) + + +def _derived_half_thresholds( + requirement: Any, + config: LmdbDecodeConfig, +) -> np.ndarray: + """Resolve the per-type half-thresholds a derived margin measures against. + + Parameters + ---------- + requirement : Any + The data requirement, whose default carries the filter radius and + whose length scale names the per-element radii it multiplies. + config : LmdbDecodeConfig + Decoder state, which carries the model type map. + + Returns + ------- + np.ndarray + Half-thresholds in Å with shape ``(ntypes,)``, read-only because every + frame of a decoder shares one table. + """ + return requirement_half_thresholds( + float(_requirement_value(requirement, "default", 0.0)), + str(_requirement_value(requirement, "length_scale", "absolute")), + None if config.type_map is None else tuple(config.type_map), + config.ntypes, + ) + + def _resolve_derived_requirement( frame: dict[str, Any], key: str, @@ -852,7 +944,7 @@ def _resolve_derived_requirement( config: LmdbDecodeConfig, ) -> None: """Resolve one derived field from normalized structural frame data.""" - if key != "min_pair_dist": + if key != "pair_margin": raise ValueError(f"Unsupported derived LMDB field {key!r}") coord = frame.get("coord") @@ -865,9 +957,16 @@ def _resolve_derived_requirement( box = frame.get("box") if box is not None and np.allclose(box, 0.0): box = None - threshold = float(_requirement_value(requirement, "default", 0.0)) frame[key] = np.array( - [compute_min_pair_dist_single(coord, box, atype, stop_below=threshold)], + [ + compute_min_pair_margin_single( + coord, + box, + atype, + _derived_half_thresholds(requirement, config), + screened=True, + ) + ], dtype=_resolve_frame_dtype(config, key), ) frame[f"find_{key}"] = np.float32(1.0) @@ -879,6 +978,7 @@ def decode_lmdb_frame( config: LmdbDecodeConfig, *, copy_arrays: bool, + derive: bool = True, ) -> dict[str, Any]: """Decode and normalize one LMDB record. @@ -894,6 +994,11 @@ def decode_lmdb_frame( Whether encoded arrays are copied while unpacking. Batch decoding sets this to ``False`` because every value is copied exactly once into its preallocated batch destination. + derive + Whether the derived data requirements are resolved on this frame. + Batch decoding sets this to ``False`` and resolves them once on the + assembled batch, where the geometric scan they need runs over all + frames at once. Returns ------- @@ -959,7 +1064,7 @@ def decode_lmdb_frame( requirements = config.data_requirements for key, requirement in requirements.items(): - if _requirement_source_policy(requirement) == "derived": + if derive and _requirement_source_policy(requirement) == "derived": _resolve_derived_requirement(frame, key, requirement, config) structural_keys = frozenset( @@ -1302,6 +1407,7 @@ def decode_lmdb_batch( int(original_key), config, copy_arrays=False, + derive=False, ) frame_nloc = frame["coord"].shape[0] if layout is None: @@ -1380,6 +1486,9 @@ def decode_lmdb_batch( batch.pop(field, None) if layout.ragged: batch["n_node"] = layout.n_node + for key, requirement in config.data_requirements.items(): + if _requirement_source_policy(requirement) == "derived": + _resolve_derived_batch(batch, key, requirement, config, layout) batch["sid"] = np.asarray([0], dtype=np.int64) return batch @@ -1601,7 +1710,7 @@ class LmdbBatchIterator: reader LMDB reader that owns metadata and the synchronous transaction. sampler - Finite iterator yielding same-nloc dataset-index batches. + Finite sized iterable yielding dataset-index batches. num_workers Decoder process count. Zero or one selects synchronous decoding. """ @@ -1617,6 +1726,8 @@ def __init__( self._reader = reader self._sampler = sampler self._epoch = 0 + self._queued_pass_length = 0 + self._last_pass_length: int | None = None self._iterator = self._iter_epoch() self._num_workers = num_workers self._pool: _LmdbPoolEntry | None = None @@ -1638,9 +1749,12 @@ def __next__(self) -> dict[str, Any]: indices, self._deferred_indices = self._deferred_indices, None batch = self._reader.decode_batch(indices) else: - batch = self._decode(self._next_indices()) + indices, self._queued_pass_length = self._next_indices() + batch = self._decode(indices) - self._schedule(self._next_indices()) + self._last_pass_length = self._queued_pass_length + indices, self._queued_pass_length = self._next_indices() + self._schedule(indices) return batch def _decode(self, indices: list[int]) -> dict[str, Any]: @@ -1719,7 +1833,7 @@ def _lose_the_pool(self, futures: "Sequence[Future[dict[str, Any]]]" = ()) -> No "decoders outlive the session that started it." ) - def _next_indices(self) -> list[int]: + def _next_indices(self) -> tuple[list[int], int]: """Return the next sampler batch and restart after exhaustion.""" try: return next(self._iterator) @@ -1728,12 +1842,17 @@ def _next_indices(self) -> list[int]: self._iterator = self._iter_epoch() return next(self._iterator) - def _iter_epoch(self) -> Iterator[list[int]]: - """Create a sampler iterator for the current epoch.""" + def _iter_epoch(self) -> Iterator[tuple[list[int], int]]: + """Tag batches with the length of the sampler pass that owns them.""" set_epoch = getattr(self._sampler, "set_epoch", None) if callable(set_epoch): set_epoch(self._epoch) - return iter(self._sampler) + # Samplers can advance before their first yield, and prefetch can + # enter another pass before the current batch reaches its consumer. + # Snapshot the length lazily, after data requirements and layout are set. + pass_length = len(self._sampler) + for indices in self._sampler: + yield indices, pass_length def _submit(self, indices: list[int]) -> list[Future[dict[str, Any]]]: """Submit one batch as balanced contiguous chunks. @@ -1774,6 +1893,13 @@ def _schedule(self, indices: list[int]) -> None: self._pending = _PendingBatch(indices, futures) if futures else None self._deferred_indices = None if futures else indices + @property + def batch_pass_length(self) -> int: + """Length of the last returned batch's pass, or the first pending pass.""" + if self._last_pass_length is None: + return len(self._sampler) + return self._last_pass_length + @property def started(self) -> bool: """Whether this iterator has acquired the shared process pool.""" @@ -2176,6 +2302,7 @@ def __init__( type_remap=self._type_remap, data_requirements=self._data_requirements, dataset=self.lmdb_path, + type_map=self._type_map, ) # Which fields carry an atom axis follows from the registered # requirements, so this cache is invalidated when they change. @@ -3826,6 +3953,7 @@ def __init__( type_remap=self._type_remap, data_requirements=self._requirements, dataset=self.lmdb_path, + type_map=self._type_map, ) # Detect PBC from the first retained frame. diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index da6b890d13..9d54519e9d 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -87,6 +87,9 @@ from deepmd.utils.compat import ( update_deepmd_input, ) +from deepmd.utils.data import ( + min_pair_dist_requirement, +) from deepmd.utils.data_system import ( get_data, process_systems, @@ -615,12 +618,23 @@ def change_bias( ) data_requirement = mock_loss.label_requirement data_requirement += training.get_additional_data_requirement(model_to_change) + # A bridged model draws its statistics under the frame filter of its + # window, as training does; no other filter applies here. + pair_filter = min_pair_dist_requirement( + model_params + if not multi_task + else model_params["model_dict"][model_branch], + 0.0, + ) + if pair_filter is not None: + data_requirement.append(pair_filter) data_single.add_data_requirement(data_requirement) nbatches = numb_batch if numb_batch != 0 else float("inf") sampled_data = make_stat_input( data_single.systems, data_single.dataloaders, nbatches, + min_pair_dist=0.0 if pair_filter is None else float(pair_filter.default), ) updated_model = training.model_change_out_bias( model_to_change, sampled_data, _bias_adjust_mode=bias_adjust_mode diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py index 30d8ab3db0..a8864f95cc 100644 --- a/deepmd/pt/model/descriptor/hybrid.py +++ b/deepmd/pt/model/descriptor/hybrid.py @@ -64,6 +64,16 @@ def __init__( ) else: raise NotImplementedError + # A bridged child fades its fitting output through a readout gate, + # which a concatenated descriptor has no fitting output to apply to. + if any( + getattr(descrpt, "bridging_switch", None) is not None + for descrpt in formatted_descript_list + ): + raise NotImplementedError( + "A hybrid descriptor cannot carry a bridged child; bridge the " + "model with a single descriptor." + ) self.descrpt_list = torch.nn.ModuleList(formatted_descript_list) self.numb_descrpt = len(self.descrpt_list) for ii in range(1, self.numb_descrpt): diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py index 9e565bfb40..7a3373da5e 100644 --- a/deepmd/pt/model/descriptor/sezm.py +++ b/deepmd/pt/model/descriptor/sezm.py @@ -69,6 +69,15 @@ cuda_infer_level, use_amp_infer, ) +from deepmd.utils.bridging import ( + BRIDGING_RECORD_VERSION, + check_bridging_record_version, + check_window_inside_cutoff, + migrate_inner_clamp_keys, +) +from deepmd.utils.element_radii import ( + contact_radius_table, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -423,12 +432,21 @@ class DescrptSeZM(BaseDescriptor, nn.Module): Random seed(s). type_map Type names. - inner_clamp_r_inner - Inner radius for distance saturation in Å. If both inner and outer radii - are set, the descriptor freezes short-range descriptor geometry inside - the zone-bridging window. - inner_clamp_r_outer - Outer radius for distance saturation in Å. + inner_clamp_f_inner + Inner radius of the zone-bridging window, as a fraction of each pair's + own length scale. When both fractions are set, the descriptor freezes + short-range geometry inside the window. Given together with + ``inner_clamp_f_outer`` or not at all. + inner_clamp_f_outer + Outer radius of the zone-bridging window, as a fraction of each pair's + own length scale. At or beyond it an edge contributes exactly as it + does without a window. + inner_clamp_scale + Length scale the two fractions measure against. ``"covalent"`` sizes + every pair by the sum of the covalent radii of its two elements, so one + window shape serves every element combination and a ``type_map`` is + required; ``"absolute"`` gives every element the unit radius, which + makes the pair scale 1 Å and the two fractions radii in Å. add_chg_spin_ebd If True, add frame-level charge/spin condition embedding to scalar type features before edge features are built. @@ -446,7 +464,7 @@ class DescrptSeZM(BaseDescriptor, nn.Module): """ _ENV_DIM: int = 1 # Use se_r style (radial only) for EnvMatStatSe compatibility - LATEST_VERSION: float = 1.2 + LATEST_VERSION: float = 1.3 def __init__( self, @@ -509,12 +527,12 @@ def __init__( trainable: bool = True, seed: int | list[int] | None = None, type_map: list[str] | None = None, - inner_clamp_r_inner: float | None = None, - inner_clamp_r_outer: float | None = None, + inner_clamp_f_inner: float | None = None, + inner_clamp_f_outer: float | None = None, + inner_clamp_scale: str = "covalent", add_chg_spin_ebd: bool = False, default_chg_spin: list[float] | None = None, use_spin: list[bool] | None = None, - **kwargs: Any, ) -> None: super().__init__() @@ -678,33 +696,60 @@ def __init__( # === Zone bridging: InnerClamp + Source Freeze Propagation Gate === # Both the geometry clamp (``InnerClamp``) and the message-passing # switch (``BridgingSwitch``) are activated together on the same - # ``[r_inner, r_outer]`` window. The clamp freezes scalar distance - # on every ``(j, k)`` edge with ``r_{jk} < r_inner``; the switch - # feeds a per-edge C3 amplitude into ``compute_edge_src_gate`` so - # that any node with a frozen neighbor cannot propagate - # information through the GNN, closing the direction / multi-hop - # leakage channels that a pure ``InnerClamp`` cannot reach. Both - # modules are parameter-free, so enabling bridging does not add - # any keys to the descriptor's state dict. - self.inner_clamp_r_inner = ( - float(inner_clamp_r_inner) if inner_clamp_r_inner is not None else None + # window, two fractions of the length scale each pair contributes + # through ``contact_radius``. The switch spans the whole window; the + # clamp freezes the distance the descriptor sees on every ``(j, k)`` + # edge two fifths into the window, just below the frame-filter point + # at the midpoint, so that the displayed distance still moves where + # the training data end and the network can fit the labels there. The + # switch feeds a per-edge C3 amplitude of the true length into + # ``compute_source_gates``, whose per-node product over every pair + # other than the edge's own is folded into the envelope of every edge + # the node emits, so that a node with a frozen neighbor is absent + # from every reduction of its other neighbors as if those edges were + # deleted, while the frozen pair keeps seeing itself at the clamped + # distance. The full per-node product also fades the learned atomic + # energy in the model, so a frozen atom contributes the bias of its + # type alone. Both modules are parameter-free, so enabling bridging + # does not add any keys to the descriptor's state dict. + if (inner_clamp_f_inner is None) != (inner_clamp_f_outer is None): + raise ValueError( + "`inner_clamp_f_inner` and `inner_clamp_f_outer` must be given " + "together." + ) + if inner_clamp_scale not in ("covalent", "absolute"): + raise ValueError( + "`inner_clamp_scale` must be 'covalent' or 'absolute', got " + f"{inner_clamp_scale!r}." + ) + self.bridging_f_inner = ( + None if inner_clamp_f_inner is None else float(inner_clamp_f_inner) ) - self.inner_clamp_r_outer = ( - float(inner_clamp_r_outer) if inner_clamp_r_outer is not None else None + self.bridging_f_outer = ( + None if inner_clamp_f_outer is None else float(inner_clamp_f_outer) ) - if ( - self.inner_clamp_r_inner is not None - and self.inner_clamp_r_outer is not None - ): - self.inner_clamp: InnerClamp | None = InnerClamp( - self.inner_clamp_r_inner, self.inner_clamp_r_outer + self.bridging_scale = str(inner_clamp_scale) + if self.bridging_f_inner is not None: + self.bridging_clamp: InnerClamp | None = InnerClamp( + self.bridging_f_inner, self.bridging_f_outer ) self.bridging_switch: BridgingSwitch | None = BridgingSwitch( - self.inner_clamp_r_inner, self.inner_clamp_r_outer + self.bridging_f_inner, self.bridging_f_outer ) else: - self.inner_clamp = None + self.bridging_clamp = None self.bridging_switch = None + # The table sizes a window, so an unbridged descriptor needs none and + # is not asked for the element symbols the covalent scale reads. It is + # derived from the configuration rather than trained, hence transient. + contact_radius = None + if self.bridging_f_inner is not None: + table = contact_radius_table( + self.ntypes, self.type_map, self.bridging_scale + ) + check_window_inside_cutoff(table, self.bridging_f_outer, self.rcut) + contact_radius = torch.tensor(table, dtype=torch.float64, device=env.DEVICE) + self.register_buffer("contact_radius", contact_radius, persistent=False) # === Env seed parameters === self.use_env_seed = bool(use_env_seed) @@ -1069,16 +1114,11 @@ def __init__( self.blocks = nn.ModuleList(blocks) # The fused convolution paths consume only the three structural rows of - # each Wigner degree block. Source-gated attention bypasses that fused - # convolution, so its dense per-edge rotations remain available. - self._wigner_free_conv = ( - self.bridging_switch is None - and bool(self.blocks) - and all( - getattr(block.so2_conv, "_cuda_conv_fn", None) is not None - and not block.so2_conv._cuda_conv_fn._compete - for block in self.blocks - ) + # each Wigner degree block. + self._wigner_free_conv = bool(self.blocks) and all( + getattr(block.so2_conv, "_cuda_conv_fn", None) is not None + and not block.so2_conv._cuda_conv_fn._compete + for block in self.blocks ) self._packed_wigner_train = bool(self.blocks) and all( getattr(block.so2_conv, "_cuda_value_train", None) is not None @@ -1184,6 +1224,37 @@ def __init__( persistent=True, ) + def pair_contact( + self, + center_type: torch.Tensor, + neighbor_type: torch.Tensor, + dtype: torch.dtype, + ) -> torch.Tensor: + """ + Gather the length scale of each edge from the types at its ends. + + Parameters + ---------- + center_type : torch.Tensor + Type index of the center atom of every edge, with shape (E,). + neighbor_type : torch.Tensor + Type index of the neighbor atom of every edge, with shape (E,). + dtype : torch.dtype + Floating-point dtype of the result. The edge cache divides the edge + distances by that result, so this is the dtype the geometry is + reduced in rather than the dtype of the incoming coordinates. + + Returns + ------- + torch.Tensor + Contact distance of each edge in Å, with shape (E, 1). + """ + radius = self.contact_radius.to(dtype=dtype) + contact = radius.index_select(0, center_type) + radius.index_select( + 0, neighbor_type + ) + return contact.unsqueeze(-1) + def forward( self, extended_coord: torch.Tensor, @@ -1206,7 +1277,12 @@ def forward( torch.Tensor, ]: """ - Compute the descriptor. + Compute the descriptor through the common descriptor interface. + + The interface returns the common five elements, so the readout gate of + a bridged descriptor does not travel through it; the SeZM model + evaluates its energies through :meth:`forward_with_edges`, which hands + the gate over. Parameters ---------- @@ -1268,7 +1344,7 @@ def forward( dtype=extended_coord.dtype, device=extended_coord.device, ) - descriptor, _, _ = self.forward_with_edges( + descriptor, _, _, _ = self.forward_with_edges( extended_coord=extended_coord, extended_atype=extended_atype, edge_index=edge_index, @@ -1481,7 +1557,7 @@ def forward_with_edges( comm_dict: dict[str, torch.Tensor] | None = None, nloc: int | None = None, vacuum_conditions: dict[str, torch.Tensor] | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: """ Compute the descriptor from a sparse edge list. @@ -1533,11 +1609,13 @@ def forward_with_edges( Returns ------- - tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] + tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None] The scalar descriptor with shape ``(nf, nloc, channels)``, the - final equivariant latent with shape ``(nf * nloc, D_final, 1, channels)`` - and, with ``vacuum_conditions``, the vacuum descriptor with shape - ``(ntypes, channels)``; ``None`` otherwise. + final equivariant latent with shape ``(nf * nloc, D_final, 1, channels)``, + with ``vacuum_conditions`` the vacuum descriptor with shape + ``(ntypes, channels)`` (``None`` otherwise), and for a bridged + descriptor the per-atom source gate with shape ``(nf, nloc, 1)`` + that fades the learned atomic energy (``None`` otherwise). """ # === Step 1. Setup dimensions === # ``n_per_frame`` is the per-frame node count: ``nloc`` in the @@ -1627,6 +1705,18 @@ def forward_with_edges( self._gate_partial_exchange, comm_dict=comm_dict ) # === Step 3. Build edge cache once (sparse edges) === + # The window of every edge is sized by the two types at its ends, so + # the descriptor, which owns the radius table, resolves it here and the + # edge cache receives one per-edge length scale. + edge_contact = ( + None + if self.contact_radius is None + else self.pair_contact( + atype_flat.index_select(0, edge_index[1]), + atype_flat.index_select(0, edge_index[0]), + self.compute_dtype, + ) + ) with nvtx_range("build_edge_cache"): edge_cache = build_edge_cache_from_edges( type_ebed=type_ebed, @@ -1639,8 +1729,9 @@ def forward_with_edges( deg_norm_floor=( self.deg_norm_floor if self.version >= 1.1 else self.eps ), - inner_clamp=self.inner_clamp, + bridging_clamp=self.bridging_clamp, bridging_switch=self.bridging_switch, + edge_contact=edge_contact, edge_envelope=self.edge_envelope, radial_basis=self.radial_basis, fused_radial=(None if self.training else self._cuda_radial_fn), @@ -1779,7 +1870,14 @@ def forward_with_edges( if vacuum_ref else None ) - return descriptor, latent.contiguous(), vacuum + node_gate = ( + None + if edge_cache.node_gate is None + else edge_cache.node_gate[:n_out_nodes] + .reshape(nf, out_nloc, 1) + .to(dtype=env.GLOBAL_PT_FLOAT_PRECISION) + ) + return descriptor, latent.contiguous(), vacuum, node_gate def _forward_blocks( self, @@ -2773,8 +2871,9 @@ def serialize(self) -> dict[str, Any]: "eps": self.eps, "trainable": self.trainable, "seed": self.seed, - "inner_clamp_r_inner": self.inner_clamp_r_inner, - "inner_clamp_r_outer": self.inner_clamp_r_outer, + "inner_clamp_f_inner": self.bridging_f_inner, + "inner_clamp_f_outer": self.bridging_f_outer, + "inner_clamp_scale": self.bridging_scale, "add_chg_spin_ebd": self.add_chg_spin_ebd, "default_chg_spin": self.default_chg_spin, "use_spin": self.use_spin, @@ -2798,6 +2897,7 @@ def deserialize(cls, data: dict[str, Any]) -> DescrptSeZM: variables = data.pop("@variables") data.pop("env_mat", None) config.pop("s2_grid_resolution", None) + migrate_inner_clamp_keys(config) obj = cls(**config) template = obj.state_dict() state = { @@ -2859,9 +2959,11 @@ def _migrate_variables( to a module: ``load_state_dict`` restores a module's own buffers before descending into its children, so a migration applied to live attributes would rewrite values the child load is about to - overwrite. Only representations are upgraded here; a difference no - rewrite can absorb stays a forward-time branch on :attr:`version`, - so a migrated descriptor never changes its own math. + overwrite. Both checkpoints and portable records pass through this + migration. A bridged receiving descriptor rejects legacy states before + its tensors are restored. Other differences that no rewrite + can absorb stay forward-time branches on :attr:`version`, so a + migrated descriptor never changes its own math. Version 1.2 moved the env-seed spin gate from the spin coordinate to the resulting environment quadratic form. For an active-spin model, @@ -2872,6 +2974,14 @@ def _migrate_variables( be activated by fine-tuning. Versions below 1.1 predate the native-spin route and retain their original forward semantics. + Version 1.3 changes only bridging. A version-1.2 state loaded into an + unbridged descriptor adopts the current record version without changing + its function. A bridged receiving descriptor rejects every pre-1.3 state: + state dictionaries do not identify the source's bridging configuration. + Fine-tuning first loads a compatible plain checkpoint into its original + unbridged model, migrates it, then transfers the state to the bridged + target. Legacy bridged checkpoints cannot be migrated this way. + Parameters ---------- variables @@ -2886,24 +2996,29 @@ def _migrate_variables( float Version the state expresses after migration. """ - if not 1.1 <= version < 1.2: - return version - - gate_key = prefix + "env_seed_embedding.spin_scale" - if self.use_spin is not None and not any(self.use_spin): - dormant_keys = ( - "spin_embedding.mag_layer2.matrix", - "spin_embedding.adam_spin_vec_weight", - "spin_embedding.adam_spin_nbr_weight", - "env_seed_embedding.spin_scale", - ) - for name in dormant_keys: - key = prefix + name - if key in variables: - variables[key] = torch.zeros_like(variables[key]) - elif gate_key in variables: - variables[gate_key] = variables[gate_key] ** 2 - return 1.2 + check_bridging_record_version( + {"inner_clamp_f_inner": self.bridging_f_inner}, version + ) + if 1.1 <= version < 1.2: + gate_key = prefix + "env_seed_embedding.spin_scale" + if self.use_spin is not None and not any(self.use_spin): + dormant_keys = ( + "spin_embedding.mag_layer2.matrix", + "spin_embedding.adam_spin_vec_weight", + "spin_embedding.adam_spin_nbr_weight", + "env_seed_embedding.spin_scale", + ) + for name in dormant_keys: + key = prefix + name + if key in variables: + variables[key] = torch.zeros_like(variables[key]) + elif gate_key in variables: + variables[gate_key] = variables[gate_key] ** 2 + version = 1.2 + + if 1.2 <= version < BRIDGING_RECORD_VERSION: + return BRIDGING_RECORD_VERSION + return version def _load_from_state_dict( self, diff --git a/deepmd/pt/model/descriptor/sezm_nn/__init__.py b/deepmd/pt/model/descriptor/sezm_nn/__init__.py index aae88f5f86..d185016cde 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/__init__.py +++ b/deepmd/pt/model/descriptor/sezm_nn/__init__.py @@ -36,7 +36,7 @@ build_edge_cache, build_edge_cache_from_edges, build_edge_type_feat, - compute_edge_src_gate, + compute_source_gates, edge_cache_to_dtype, ) from .embedding import ( @@ -190,7 +190,7 @@ "build_m_major_l_index", "build_merged_state_dict", "build_rotate_inv_rescale", - "compute_edge_src_gate", + "compute_source_gates", "edge_cache_to_dtype", "fold_lora_state_dict_keys", "get_promoted_dtype", diff --git a/deepmd/pt/model/descriptor/sezm_nn/attention.py b/deepmd/pt/model/descriptor/sezm_nn/attention.py index ec20c4ac93..5b889b0f37 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/attention.py +++ b/deepmd/pt/model/descriptor/sezm_nn/attention.py @@ -22,11 +22,17 @@ def segment_envelope_gated_softmax( n_nodes: int, z_bias_raw: torch.Tensor, eps: float, - src_weight: torch.Tensor | None = None, ) -> torch.Tensor: """ Compute destination-wise envelope-gated softmax attention. + The physical mass of an edge is ``edge_env**2 * exp(logits)`` and the + denominator of its destination is the sum of the edge masses plus the + positive null mass ``softplus(z_bias_raw) + eps``. An edge whose envelope + is zero is therefore absent from both, which is what the frozen-zone + invariance of a bridged model needs, since its envelope carries the + source gate. + Parameters ---------- logits @@ -42,16 +48,6 @@ def segment_envelope_gated_softmax( Softplus is applied to keep the bias strictly positive. eps Small positive floor added to the physical null mass. - src_weight - Optional per-edge source-side multiplier with shape (E, 1) or - (E,). When provided, the physical per-edge mass is - ``edge_env**2 * src_weight * exp(logits)`` and the denominator is the - sum of edge masses plus the positive null mass - ``softplus(z_bias_raw) + eps``. - ``src_weight = 0`` therefore removes the source from both the - numerator and the denominator, which is what SFPG needs so that - a muted source does not even leak through the softmax - normalization. Returns ------- @@ -65,30 +61,16 @@ def segment_envelope_gated_softmax( torch.float32 if input_dtype in (torch.float16, torch.bfloat16) else input_dtype ) - # === Step 1. Build factor-wise effective logits === - # Computing the logarithms before multiplying the factors avoids losing a - # physically nonzero edge when ``edge_env**2 * src_weight`` underflows. + # === Step 1. Build the effective logits === + # Taking the logarithm of the envelope before adding it to the logits + # avoids losing a physically nonzero edge when ``edge_env**2`` underflows. logits_2d = logits.reshape(n_edge, n_channel).to(dtype=compute_dtype) edge_env_1d = edge_env.reshape(n_edge).to(dtype=compute_dtype) edge_positive = edge_env_1d > 0.0 ones = torch.ones_like(edge_env_1d) log_weight = 2.0 * torch.log(torch.where(edge_positive, edge_env_1d, ones)) - active = edge_positive - source_ratio: torch.Tensor | None = None - if src_weight is not None: - source_weight = src_weight.reshape(n_edge).to(dtype=compute_dtype) - source_positive = source_weight > 0.0 - safe_source = torch.where(source_positive, source_weight, ones) - source_scale = safe_source.detach() - log_weight = log_weight + torch.log(source_scale) - source_ratio = torch.where( - source_positive, - source_weight / source_scale, - torch.zeros_like(source_weight), - ) - active = active & source_positive effective_logits = torch.where( - active.reshape(n_edge, 1), + edge_positive.reshape(n_edge, 1), logits_2d + log_weight.reshape(n_edge, 1), torch.full_like(logits_2d, float("-inf")), ) @@ -115,11 +97,6 @@ def segment_envelope_gated_softmax( # === Step 3. Normalize edge and null masses in the shared shifted frame === edge_exp = torch.exp(effective_logits - edge_max) - if source_ratio is not None: - # ``source_scale`` carries the forward magnitude in log space, while - # this linear ratio carries derivatives without differentiating - # ``log(src_weight)`` at extremely small positive gates. - edge_exp = edge_exp * source_ratio.reshape(n_edge, 1) denom_sum = torch.zeros( n_nodes, n_channel, diff --git a/deepmd/pt/model/descriptor/sezm_nn/dens.py b/deepmd/pt/model/descriptor/sezm_nn/dens.py index 10cbc2fb86..547b08d6d4 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/dens.py +++ b/deepmd/pt/model/descriptor/sezm_nn/dens.py @@ -669,6 +669,7 @@ def forward( aparam: torch.Tensor | None = None, vacuum_descriptor: torch.Tensor | None = None, return_components: bool = False, + node_gate: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """ Run the parallel `dens` fitting branches. @@ -692,6 +693,9 @@ def forward( `(ntypes, dim_descrpt)`, required by ``vacuum_ref``. return_components If true, also return the clean-force and denoising branches. + node_gate + Per-atom source gate of a bridged descriptor with shape + `(nf, nloc, 1)`; the learned part of the energy fades with it. Returns ------- @@ -709,6 +713,7 @@ def forward( fparam=fparam, aparam=aparam, vacuum_descriptor=vacuum_descriptor, + node_gate=node_gate, ) clean_force = self.direct_force_head(latent).view(nf, nloc, 3) denoising_force = self.denoising_head(latent).view(nf, nloc, 3) diff --git a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py index 1953c2f620..4a67e2d1f9 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py +++ b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py @@ -65,7 +65,20 @@ class EdgeFeatureCache(NamedTuple): Radial basis with shape (E, n_radial). The C^3 cutoff envelope is already baked in. edge_env - C^3 cutoff envelope weights with shape (E, 1). + C^3 cutoff envelope weights with shape (E, 1). In bridging mode the + Source Freeze Propagation Gate is folded into it: the envelope of the + edge ``j -> k`` is multiplied by the product of the switching + amplitudes ``w(r_{jl})`` of every pair of the source ``j`` other than + the pair ``(j, k)`` itself, so an edge whose source has a third + neighbor in the frozen zone vanishes everywhere the envelope enters + (messages, environment seed, attention masses and degree), exactly as + an edge beyond the cutoff does, while the edge between the two atoms + of a frozen pair keeps its clamped geometry. + node_gate + Per-node source gate ``eta[j] = prod_{k in N(j)} w(r_{jk})`` with + shape (N,), the product over every pair of the node; ``None`` without + bridging. The atomic model fades the learned atomic energy of node + ``j`` with it. deg Envelope-squared smooth degree with shape (N,), computed as ``sum(edge_env**2)`` over incoming edges. @@ -90,17 +103,6 @@ class EdgeFeatureCache(NamedTuple): Lazy cache for endpoint CSR views used by segmented accelerated operators, keyed by endpoint role (``"dst"`` or ``"src"``). Built once per step and shared by every consumer. - edge_src_gate - Optional per-edge Source Freeze Propagation Gate (SFPG) weight with - shape (E, 1). Equals ``eta[src]`` where - ``eta[j] = prod_{k in N(j)} w(r_{jk})`` and ``w`` is the - :class:`BridgingSwitch` C3 switching amplitude. Present only when - the model runs in bridging mode; ``None`` otherwise. Aggregation - sites (``GeometricInitialEmbedding``, ``EnvironmentInitialEmbedding``, - ``SO2Convolution``) multiply their per-edge message contribution - by this gate to forbid any node whose local neighborhood enters - the frozen zone from propagating information along its outgoing - edges. """ src: torch.Tensor @@ -116,8 +118,8 @@ class EdgeFeatureCache(NamedTuple): D_to_m_cache: dict[str, torch.Tensor] | None = None Dt_from_m_cache: dict[str, torch.Tensor] | None = None csr_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] | None = None - edge_src_gate: torch.Tensor | None = None edge_quat: torch.Tensor | None = None + node_gate: torch.Tensor | None = None def cached_edge_csr( @@ -161,30 +163,39 @@ def cached_edge_csr( return order, row_ptr -def compute_edge_src_gate( +def compute_source_gates( *, edge_len: torch.Tensor, + edge_contact: torch.Tensor, src: torch.Tensor, n_nodes: int, - bridging_switch: Callable[[torch.Tensor], torch.Tensor], + bridging_switch: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], edge_keep_f: torch.Tensor | None = None, node_partial_exchange: Callable[[torch.Tensor], torch.Tensor] | None = None, -) -> torch.Tensor: +) -> tuple[torch.Tensor, torch.Tensor]: """ - Compute the per-edge source gate for SFPG from edge lengths. - - The gate implements a per-node "non-frozen confidence" and broadcasts - it back to edges along the source axis:: - - w_e = bridging_switch(edge_len_e) in [0, 1] - eta_j = prod_{e: src_e = j} w_e in [0, 1] - gate_e = eta_{src_e} in [0, 1] - - ``w_e = 0`` at ``r_{jk} <= r_inner`` ensures ``eta_j = 0`` for any - node with at least one neighbor in the frozen zone. Masked edges - (padding, excluded type pairs) must contribute the multiplicative - identity ``1`` so they never spuriously mute a valid source node; - callers supply ``edge_keep_f`` for this. + Compute the per-node and per-edge source gates of SFPG from edge lengths. + + The per-node gate is the "non-frozen confidence" of a node, the per-edge + gate the same product with the edge's own pair left out:: + + w_e = bridging_switch(edge_len_e, edge_contact_e) in [0, 1] + eta_j = prod_{e: src_e = j} w_e in [0, 1] + gate_e = eta_{src_e} / w_e = prod_{e': src_e' = src_e, e' != e} w_e' + + ``w_e = 0`` inside the inner radius of that edge's own pair ensures + ``eta_j = 0`` for any node with at least one neighbor in the frozen zone, + and ``gate_e = 0`` for every edge such a node emits except the edge to + that frozen neighbor. The cache builder multiplies ``gate_e`` into the + edge envelope, which every consumer weights by, so a muted edge is + indistinguishable from a deleted one: the frozen atom disappears from the + environment of every third atom, while the two atoms of a frozen pair + keep seeing each other at the clamped distance. The fitting scales the + deviation of its output from the bias of the type by ``eta_j``, so the + frozen atom contributes that bias alone. Masked edges (padding, + excluded type pairs) must contribute the multiplicative identity ``1`` so + they never spuriously mute a valid source node; callers supply + ``edge_keep_f`` for this. The product is **not** realised by ``scatter_reduce(reduce="prod")``: its registered backward handles exact zeros with a data-dependent @@ -193,12 +204,14 @@ def compute_edge_src_gate( path's double-backward tracing. Instead, the product is decomposed into a log-sum on non-zero contributions combined with an explicit "any zero per group" indicator that routes the frozen case through - ``torch.where``. Both branches use only shape-preserving standard - ops (``scatter_add``, ``where``, ``exp``, ``log``) with backed - symints, so the graph survives symbolic tracing cleanly. + ``torch.where``; the leave-one-out product subtracts the edge's own + log-amplitude and zero flag from the totals of its source. Both branches + use only shape-preserving standard ops (``scatter_add``, ``where``, + ``exp``, ``log``) with backed symints, so the graph survives symbolic + tracing cleanly. The gradient consequence at the plateau is exact: ``BridgingSwitch`` - places ``w'(r) = 0`` for every ``r <= r_inner``, so the chain rule + places ``w'(r) = 0`` everywhere inside the inner radius, so the chain rule ``d eta / d r = (leave-one-out factor) * w'(r) = anything * 0 = 0`` holds regardless of how the muted ``torch.where`` branch treats the upstream gradient. In the transition zone every edge has strictly @@ -209,13 +222,17 @@ def compute_edge_src_gate( ---------- edge_len Per-edge distances with shape (E, 1). + edge_contact + Per-edge length scale with shape (E, 1) in Å, the sum of the two + endpoint radii. The switch measures the window against it, so every + element combination gets its own pair of radii. src Source node indices with shape (E,). n_nodes Total number of nodes N. bridging_switch - Callable ``r -> w(r)`` with ``w: [0, ∞) -> [0, 1]``, typically a - :class:`BridgingSwitch` instance. + Callable ``(r, contact) -> w`` with ``w: [0, ∞) -> [0, 1]``, typically + a :class:`BridgingSwitch` instance. edge_keep_f Optional per-edge keep weights with shape (E, 1), with ``0`` on masked edges and ``1`` on kept edges. If provided, masked edges @@ -231,12 +248,14 @@ def compute_edge_src_gate( Returns ------- - torch.Tensor - Per-edge source gate with shape (E, 1), aligned on the same edge - axis as the rest of the cache. + node_gate : torch.Tensor + Per-node source gate ``eta`` with shape (N,). + edge_gate : torch.Tensor + Per-edge leave-one-out gate with shape (E, 1), aligned on the same + edge axis as the rest of the cache. """ # === Step 1. Per-edge switching amplitude w(r) in [0, 1] === - edge_w = bridging_switch(edge_len) # (E, 1) + edge_w = bridging_switch(edge_len, edge_contact) # (E, 1) if edge_keep_f is not None: # Force w = 1 on masked edges so they are neutral for the product. edge_w = edge_w * edge_keep_f + (1.0 - edge_keep_f) @@ -260,9 +279,10 @@ def compute_edge_src_gate( # the hard-freeze rule. Float count (values are small integers, exact # in fp) so both partials ride ONE border-exchange tensor when # completing across ranks. + is_zero_f = is_zero.to(edge_w.dtype) zero_count = torch.zeros( n_nodes, dtype=edge_w.dtype, device=edge_w.device - ).scatter_add(0, src, is_zero.to(edge_w.dtype)) + ).scatter_add(0, src, is_zero_f) # === Step 3b. Cross-rank completion of the per-node partials === # A rank only holds edges whose dst is owned, so the src-keyed sums @@ -276,12 +296,21 @@ def compute_edge_src_gate( log_eta = packed[..., 0] zero_count = packed[..., 1] - eta_nonzero_path = torch.exp(log_eta) - any_zero = zero_count > 0.5 - - # === Step 4. Combine and broadcast back to edges via source === - eta = torch.where(any_zero, torch.zeros_like(eta_nonzero_path), eta_nonzero_path) - return eta.index_select(0, src).unsqueeze(-1) + # === Step 4. Per-node gate === + node_gate = torch.where( + zero_count > 0.5, torch.zeros_like(log_eta), torch.exp(log_eta) + ) # (N,) + + # === Step 5. Per-edge leave-one-out gate === + # The edge's own factor is removed from its source's totals: its + # log-amplitude (0 when it is a frozen edge, because ``safe_w`` is 1 there) + # and its zero flag, so the remaining product is over the other pairs. + loo_log = log_eta.index_select(0, src) - log_safe # (E,) + loo_zero = zero_count.index_select(0, src) - is_zero_f # (E,) + edge_gate = torch.where( + loo_zero > 0.5, torch.zeros_like(loo_log), torch.exp(loo_log) + ) + return node_gate, edge_gate.unsqueeze(-1) @torch.amp.autocast("cuda", enabled=False) @@ -443,8 +472,9 @@ def build_edge_cache_from_edges( compute_dtype: torch.dtype, eps: float, deg_norm_floor: float, - inner_clamp: Callable[[torch.Tensor], torch.Tensor] | None, - bridging_switch: Callable[[torch.Tensor], torch.Tensor] | None, + bridging_clamp: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None, + bridging_switch: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None, + edge_contact: torch.Tensor | None, edge_envelope: Callable[[torch.Tensor], torch.Tensor], radial_basis: Callable[[torch.Tensor], torch.Tensor], has_exclude_types: bool, @@ -478,15 +508,21 @@ def build_edge_cache_from_edges( deg_norm_floor Floor added to the envelope-squared degree before inverse-sqrt normalization (see :func:`_finalize_edge_cache`). - inner_clamp - Optional inner clamp used to freeze short-range geometry below `r_inner`. + bridging_clamp + Optional distance clamp that freezes the geometry the descriptor sees + for a pair inside its bridging window (see :class:`InnerClamp`). bridging_switch - Optional C3 switching amplitude ``w(r) -> [0, 1]`` that drives - the Source Freeze Propagation Gate. When provided, a per-edge - ``edge_src_gate`` is computed from the node-wise product of - ``w(r_{jk})`` along each source node's outgoing edges. Masked - edges (``edge_keep=False``) are forced to ``w=1`` so they never - leak into the product. + Optional C3 switching amplitude ``w(r, contact) -> [0, 1]`` that drives + the Source Freeze Propagation Gate. When provided, the product of + ``w(r_{jl})`` over the true lengths of every pair of the source node + other than the edge's own pair is folded into the envelope of every + edge that node emits, and the full product is stored on the cache as + ``node_gate``. Masked edges (``edge_keep=False``) are forced to + ``w=1`` so they never leak into the product. + edge_contact + Per-edge length scale with shape (E, 1) in Å that both window modules + measure their fractions against, or ``None`` when neither is given. + The descriptor owns the per-type radii and resolves it before the call. edge_envelope C^3 edge envelope module. radial_basis @@ -534,8 +570,9 @@ def build_edge_cache_from_edges( # === Step 3. Edge length, envelope, and radial basis === with nvtx_range("envelope"): edge_len = safe_norm(edge_vec, eps) - if inner_clamp is not None: - clamped = inner_clamp(edge_len) + edge_len_true = edge_len + if bridging_clamp is not None: + clamped = bridging_clamp(edge_len, edge_contact) scale = clamped / edge_len edge_vec = edge_vec * scale edge_len = clamped @@ -561,21 +598,30 @@ def build_edge_cache_from_edges( edge_type_feat = edge_type_feat * edge_keep_f.to(dtype=edge_type_feat.dtype) # === Step 6. Source Freeze Propagation Gate (optional) === - # The sparse-edge path packs masked dummy edges so the compiled graph sees - # a statically non-empty, non-singular edge tensor. ``edge_keep_f`` rewrites - # any such slot to ``w=1`` inside ``compute_edge_src_gate``, keeping the - # product reduction unaffected by padding. - edge_src_gate: torch.Tensor | None = None + # The leave-one-out gate is folded into the envelope, the one factor every + # consumer of an edge weights by: the messages, the environment seed, the + # attention masses and the degree normalization all see a muted edge + # exactly as they see an edge beyond the cutoff. The gate reads the true + # pair length, so it closes where the clamp above has frozen the geometry. + # The radial basis stays ungated because it feeds networks whose output + # the envelope multiplies. The sparse-edge path packs masked dummy edges + # so the compiled graph sees a statically non-empty, non-singular edge + # tensor; ``edge_keep_f`` rewrites any such slot to ``w=1`` inside + # ``compute_source_gates``, keeping the product reduction unaffected by + # padding. The per-node product travels on the cache to the atomic model. + node_gate = None if bridging_switch is not None: with nvtx_range("src_gate"): - edge_src_gate = compute_edge_src_gate( - edge_len=edge_len, + node_gate, edge_gate = compute_source_gates( + edge_len=edge_len_true, + edge_contact=edge_contact, src=src, n_nodes=n_nodes, bridging_switch=bridging_switch, edge_keep_f=edge_keep_f, node_partial_exchange=node_partial_exchange, ) + edge_env = edge_env * edge_gate return _finalize_edge_cache( n_nodes=n_nodes, @@ -589,7 +635,7 @@ def build_edge_cache_from_edges( Dt_full=Dt_full, edge_quat=edge_quat, deg_norm_floor=deg_norm_floor, - edge_src_gate=edge_src_gate, + node_gate=node_gate, ) @@ -667,7 +713,7 @@ def _finalize_edge_cache( Dt_full: torch.Tensor | None, edge_quat: torch.Tensor, deg_norm_floor: float, - edge_src_gate: torch.Tensor | None = None, + node_gate: torch.Tensor | None = None, ) -> EdgeFeatureCache: """ Assemble the shared `EdgeFeatureCache` layout. @@ -702,9 +748,9 @@ def _finalize_edge_cache( normalization. A tiny ``eps`` reproduces the legacy behavior; an ``O(1)`` value makes sparse-neighborhood features vanish smoothly at ``rcut`` instead of saturating and kinking. - edge_src_gate - Optional per-edge SFPG weight with shape (E, 1). ``None`` in - non-bridging mode. + node_gate + Per-node source gate with shape (N,) of a bridged descriptor, or + ``None``. Returns ------- @@ -733,8 +779,8 @@ def _finalize_edge_cache( D_to_m_cache={}, Dt_from_m_cache={}, csr_cache={}, - edge_src_gate=edge_src_gate, edge_quat=edge_quat, + node_gate=node_gate, ) @@ -871,8 +917,9 @@ def edge_cache_to_dtype( """ Convert all floating-point tensors in EdgeFeatureCache to the specified dtype. - Integer tensors (src, dst) are unchanged. This is a standalone function - (not a method) to keep it side-effect free. + Integer tensors (src, dst) are unchanged, and so is the node gate, which + no block reads: it is applied to the fitting output in global precision. + This is a standalone function (not a method) to keep it side-effect free. Parameters ---------- @@ -890,18 +937,14 @@ def edge_cache_to_dtype( # Use local variables with explicit None check and assignment. _D_full = cache.D_full _Dt_full = cache.Dt_full - _edge_src_gate = cache.edge_src_gate _edge_quat = cache.edge_quat D_full: torch.Tensor | None = None Dt_full: torch.Tensor | None = None - edge_src_gate: torch.Tensor | None = None edge_quat: torch.Tensor | None = None if _D_full is not None: D_full = _D_full.to(dtype=dtype) if _Dt_full is not None: Dt_full = _Dt_full.to(dtype=dtype) - if _edge_src_gate is not None: - edge_src_gate = _edge_src_gate.to(dtype=dtype) if _edge_quat is not None: edge_quat = _edge_quat.to(dtype=dtype) @@ -921,6 +964,6 @@ def edge_cache_to_dtype( D_to_m_cache=None if cache.D_to_m_cache is None else {}, Dt_from_m_cache=None if cache.Dt_from_m_cache is None else {}, csr_cache=None if cache.csr_cache is None else dict(cache.csr_cache), - edge_src_gate=edge_src_gate, edge_quat=edge_quat, + node_gate=cache.node_gate, ) diff --git a/deepmd/pt/model/descriptor/sezm_nn/embedding.py b/deepmd/pt/model/descriptor/sezm_nn/embedding.py index f3215f3072..296205002a 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/embedding.py +++ b/deepmd/pt/model/descriptor/sezm_nn/embedding.py @@ -251,8 +251,8 @@ def forward( Optional per-edge neighbor-spin l=1 message with shape (E, 3, C) for the native spin scheme (built by ``SpinEmbedding.edge_l1``). It is added to the l=1 rows of the per-edge message, so it shares this - module's source gate, scatter and degree normalization with the - geometric message. + module's scatter and degree normalization with the geometric + message. Returns ------- @@ -281,12 +281,11 @@ def forward( # === Step 3. Broadcast radial features per row === # Each non-scalar packed row reuses the radial feature of its degree l. - # The fused operator spans this broadcast and the scatter of Step 5, so + # The fused operator spans this broadcast and the scatter of Step 4, so # it takes over whenever nothing else joins the message in between. if ( self._can_fuse_scatter(zonal_coupling) and spin_l1_message is None - and edge_cache.edge_src_gate is None and edge_cache.csr_cache is not None ): return self.forward_fused_scatter( @@ -303,23 +302,13 @@ def forward( # === Step 3b. Fold in the neighbor-spin l=1 message (native spin) === # The l=1 coefficients occupy the first three packed non-scalar rows, so # the neighbor-spin message joins the geometric message there and then - # shares the source gate, scatter and degree normalization below. + # shares the scatter and degree normalization below. if spin_l1_message is not None: non_scalar_message = non_scalar_message.index_add( 1, self.l1_local_index, spin_l1_message ) - # === Step 4. Source Freeze Propagation Gate (optional) === - # Mute messages emitted by nodes whose local neighborhood enters - # the frozen zone. ``edge_src_gate`` is ``None`` outside bridging - # mode so this is a no-op in normal training. - src_gate = edge_cache.edge_src_gate - if src_gate is not None: - non_scalar_message = non_scalar_message * src_gate.to( - dtype=non_scalar_message.dtype - ).unsqueeze(-1) - - # === Step 5. Scatter to nodes and normalize === + # === Step 4. Scatter to nodes and normalize === # Avoid advanced-index writeback (out[:, non_scalar_row_index, :]) which produces a copy. non_scalar_out = out.new_zeros( n_nodes, self.non_scalar_row_index.numel(), self.channels @@ -699,11 +688,6 @@ def forward( # outer = r_tilde[:, :, None] * g[:, None, :] # (E, coord_dim, embed_dim) outer = torch.einsum("ei,ej->eij", r_tilde, g) # (E, coord_dim, embed_dim) outer_flat = outer.reshape(-1, self.coord_dim * self.embed_dim) - # Source Freeze Propagation Gate: mute the outer-product contribution - # of any edge whose source node has a neighbor in the frozen zone. - src_gate = edge_cache.edge_src_gate - if src_gate is not None: - outer_flat = outer_flat * src_gate.to(dtype=outer_flat.dtype) env_agg = outer_flat.new_zeros(n_nodes, self.coord_dim * self.embed_dim) env_agg.index_add_(0, dst, outer_flat) env_agg = env_agg.reshape(n_nodes, self.coord_dim, self.embed_dim) @@ -1054,9 +1038,9 @@ def edge_l1( (neighbor) spin, scaled by a per-source-type per-channel weight and gated by the C^3 envelope. The message is returned per edge; the geometric initial embedding folds it into the l=1 rows and applies the - shared source gate, scatter and degree normalization, so a neighbor's - spin direction enters an atom's l=1 backbone before any interaction - block (the spin analogue of the geometric initial embedding). + shared scatter and degree normalization, so a neighbor's spin direction + enters an atom's l=1 backbone before any interaction block (the spin + analogue of the geometric initial embedding). Parameters ---------- diff --git a/deepmd/pt/model/descriptor/sezm_nn/radial.py b/deepmd/pt/model/descriptor/sezm_nn/radial.py index b1fac93934..73fa35ff17 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/radial.py +++ b/deepmd/pt/model/descriptor/sezm_nn/radial.py @@ -274,62 +274,89 @@ def forward(self, dst: torch.Tensor) -> torch.Tensor: class InnerClamp(nn.Module): """ - C3-continuous inner distance clamping for zone bridging. - - Applies a septic Hermite polynomial transition that freezes distances - below ``r_inner`` to the constant ``r_inner``, then smoothly transitions - back to identity at ``r_outer``:: - - r̃(r) = r_inner if r <= r_inner - r̃(r) = r_inner + (r_outer - r_inner) * h(t) if r_inner < r < r_outer - r̃(r) = r if r >= r_outer - - h(t) = 20t^4 - 45t^5 + 36t^6 - 10t^7, t = (r - r_inner) / (r_outer - r_inner) - - Boundary conditions: - ``h(0)=0``, ``h(1)=1``, ``h'(0)=0``, ``h'(1)=1``, - ``h''(0)=0``, ``h''(1)=0``, ``h'''(0)=0``, ``h'''(1)=0``. - This ensures C3 continuity: ``dr̃/dr = 0`` at r_inner (frozen zone) and - ``dr̃/dr = 1`` at r_outer (identity zone), with matched second and third - derivatives at both boundaries. + C3-continuous distance clamp of the DPA4 bridging window. + + The window of a pair is two dimensionless fractions of that pair's own + length scale ``s``, the sum of the covalent radii of its two elements, so + one window shape serves every element combination. In the reduced distance + ``u = r / s`` the clamp freezes everything below the freeze point + ``f_z`` at that value and returns to the identity at ``f_outer``:: + + ũ = f_z if u <= f_z + ũ = f_z + (f_outer - f_z) * h(t) if f_z < u < f_outer + ũ = u if u >= f_outer + + h(t) = 20t^4 - 45t^5 + 36t^6 - 10t^7, t = (u - f_z) / (f_outer - f_z) + f_z = f_inner + 0.4 (f_outer - f_inner) + + and the clamped distance is ``r̃ = s * ũ``. Boundary conditions + ``h(0)=0``, ``h(1)=1``, ``h'(0)=0``, ``h'(1)=1``, ``h''(0)=h''(1)=0``, + ``h'''(0)=h'''(1)=0`` give C3 continuity: ``dr̃/dr = 0`` at the freeze + point and ``dr̃/dr = 1`` at the outer radius (identity zone), with + matched second and third derivatives at both. + + The freeze point sits two fifths into the window, one tenth of the width + below the frame-filter point at the midpoint, where the training data of a + bridged model end. The displayed distance therefore still moves at the + filter point, with slope ``h'(1/6) = 0.22``, so the network can fit the + labels just above it through the pair's own distance; frozen at the filter + point itself, the learned pair energy could only follow the switch + amplitude, whose slope at the midpoint pins the wall to a stiffness the + labels do not have. Below the filter point the displayed distance keeps + falling by ``0.6 h(1/6) = 0.62 %`` of the window width (5 mÅ for a C-C + pair) before it freezes; that band is the only range of displayed + distances no retained frame covers, and its shallowness is what bounds the + network's extrapolation there. The switch of the window + (:class:`BridgingSwitch`) spans the whole window, so below the freeze point + the learned energy of the pair follows the switch amplitude alone. Parameters ---------- - r_inner : float - Freeze radius in Å. Distances below this are clamped to ``r_inner``. - r_outer : float - Outer boundary of the transition zone in Å. Above this, ``r̃ = r``. + f_inner : float + Inner radius of the window as a fraction of the pair length scale. + f_outer : float + Outer radius as a fraction of the pair length scale. At or above it the + clamp is the identity. Raises ------ ValueError - If ``r_inner >= r_outer`` or either is non-positive. + If ``f_inner >= f_outer`` or either is non-positive. """ - def __init__(self, r_inner: float, r_outer: float) -> None: + #: Position of the freeze point inside the window, as a fraction of its width. + FREEZE_FRACTION = 0.4 + + def __init__(self, f_inner: float, f_outer: float) -> None: super().__init__() - if r_inner <= 0 or r_outer <= 0: - raise ValueError("r_inner and r_outer must be positive") - if r_inner >= r_outer: - raise ValueError(f"r_inner ({r_inner}) must be < r_outer ({r_outer})") - self.r_inner = float(r_inner) - self.r_outer = float(r_outer) + if f_inner <= 0 or f_outer <= 0: + raise ValueError("f_inner and f_outer must be positive") + if f_inner >= f_outer: + raise ValueError(f"f_inner ({f_inner}) must be < f_outer ({f_outer})") + self.f_inner = float(f_inner) + self.f_outer = float(f_outer) + self.f_freeze = self.f_inner + self.FREEZE_FRACTION * ( + self.f_outer - self.f_inner + ) - def forward(self, r: torch.Tensor) -> torch.Tensor: + def forward(self, r: torch.Tensor, contact: torch.Tensor) -> torch.Tensor: """ - Apply inner distance clamping. + Apply the distance clamp. Parameters ---------- r : torch.Tensor Pair distances with shape (...) or (..., 1) in Å. + contact : torch.Tensor + Length scale of each pair in Å, broadcastable against ``r``. Returns ------- torch.Tensor - Clamped distances r̃ with the same shape as input. + Clamped distances r̃ with the same shape as ``r``. """ - t = ((r - self.r_inner) / (self.r_outer - self.r_inner)).clamp(0.0, 1.0) + u = r / contact + t = ((u - self.f_freeze) / (self.f_outer - self.f_freeze)).clamp(0.0, 1.0) t2 = t * t t4 = t2 * t2 # h(t) = 20t^4 - 45t^5 + 36t^6 - 10t^7 @@ -339,11 +366,11 @@ def forward(self, r: torch.Tensor) -> torch.Tensor: # h''(0)=0, h''(1)=0 # h'''(0)=0, h'''(1)=0 h = t4 * (20.0 + t * (-45.0 + t * (36.0 - 10.0 * t))) - interpolated = self.r_inner + (self.r_outer - self.r_inner) * h - # Identity zone: r >= r_outer returns r directly. - # Both branches have matching first three derivatives at r_outer, + interpolated = contact * (self.f_freeze + (self.f_outer - self.f_freeze) * h) + # Identity zone: u >= f_outer returns r directly. + # Both branches have matching first three derivatives there, # so torch.where preserves C3 continuity here. - return torch.where(r >= self.r_outer, r, interpolated) + return torch.where(u >= self.f_outer, r, interpolated) class BridgingSwitch(nn.Module): @@ -354,11 +381,16 @@ class BridgingSwitch(nn.Module): that measures how far an edge sits outside the frozen zone. It is the elementary piece the Source Freeze Propagation Gate (SFPG) aggregates into a per-node "non-frozen confidence" via a product - over each source node's outgoing edges:: + over each source node's outgoing edges. + + The window of a pair is two dimensionless fractions of that pair's own + length scale ``s``, the sum of the covalent radii of its two elements, so + one window shape serves every element combination. In the reduced distance + ``u = r / s``:: - w(r) = 0 if r <= r_inner (frozen) - w(r) = h((r - r_inner) / (r_outer - r_inner)) if r_inner < r < r_outer (transition) - w(r) = 1 if r >= r_outer (normal) + w = 0 if u <= f_inner (frozen) + w = h((u - f_inner) / (f_outer - f_inner)) if f_inner < u < f_outer (transition) + w = 1 if u >= f_outer (normal) h(t) = 35 t^4 - 84 t^5 + 70 t^6 - 20 t^7 @@ -368,38 +400,40 @@ class BridgingSwitch(nn.Module): h(1)=1, h'(1) = h''(1) = h'''(1) = 0 The vanishing first three derivatives at both endpoints give - ``w \in C^3(\mathbb{R}_{\ge 0})`` with zero slope/curvature at - ``r_inner`` and ``r_outer``, so forces (first derivatives) and the + ``w \in C^3(\mathbb{R}_{\ge 0})`` with zero slope/curvature at both + radii of every pair, so forces (first derivatives) and the force derivatives consumed by second-order training stay continuous across both zone boundaries. - The surrounding infrastructure (``compute_edge_src_gate``) owns the - per-node product reduction and broadcast; this module only encodes - the scalar amplitude shape. + The surrounding infrastructure (``compute_source_gates``) owns the + per-node product reduction and the per-edge leave-one-out product; this + module only encodes the scalar amplitude shape. Parameters ---------- - r_inner : float - Inner radius in Å. At or below this distance ``w = 0``. - r_outer : float - Outer radius in Å. At or above this distance ``w = 1``. + f_inner : float + Inner radius as a fraction of the pair length scale. At or below it + ``w = 0``. + f_outer : float + Outer radius as a fraction of the pair length scale. At or above it + ``w = 1``. Raises ------ ValueError - If ``r_inner <= 0``, ``r_outer <= 0``, or ``r_inner >= r_outer``. + If ``f_inner <= 0``, ``f_outer <= 0``, or ``f_inner >= f_outer``. """ - def __init__(self, r_inner: float, r_outer: float) -> None: + def __init__(self, f_inner: float, f_outer: float) -> None: super().__init__() - if r_inner <= 0 or r_outer <= 0: - raise ValueError("r_inner and r_outer must be positive") - if r_inner >= r_outer: - raise ValueError(f"r_inner ({r_inner}) must be < r_outer ({r_outer})") - self.r_inner = float(r_inner) - self.r_outer = float(r_outer) - - def forward(self, r: torch.Tensor) -> torch.Tensor: + if f_inner <= 0 or f_outer <= 0: + raise ValueError("f_inner and f_outer must be positive") + if f_inner >= f_outer: + raise ValueError(f"f_inner ({f_inner}) must be < f_outer ({f_outer})") + self.f_inner = float(f_inner) + self.f_outer = float(f_outer) + + def forward(self, r: torch.Tensor, contact: torch.Tensor) -> torch.Tensor: """ Evaluate the C3 switching amplitude. @@ -407,19 +441,23 @@ def forward(self, r: torch.Tensor) -> torch.Tensor: ---------- r : torch.Tensor Pair distances with shape (...) or (..., 1) in Å. + contact : torch.Tensor + Length scale of each pair in Å, broadcastable against ``r``. Returns ------- torch.Tensor - Switching amplitudes in ``[0, 1]`` with the same shape as input. + Switching amplitudes in ``[0, 1]`` with the same shape as ``r``. """ - t = ((r - self.r_inner) / (self.r_outer - self.r_inner)).clamp(0.0, 1.0) + t = ((r / contact - self.f_inner) / (self.f_outer - self.f_inner)).clamp( + 0.0, 1.0 + ) t2 = t * t t4 = t2 * t2 # h(t) = 35 t^4 - 84 t^5 + 70 t^6 - 20 t^7 (Horner form). # Degree-7 smootherstep: the unique polynomial of this degree that - # hits ``w(r_inner)=0, w(r_outer)=1`` together with C3 flatness at - # both radii. + # hits ``w = 0`` at the inner radius and ``w = 1`` at the outer one + # together with C3 flatness at both. return t4 * (35.0 + t * (-84.0 + t * (70.0 - 20.0 * t))) diff --git a/deepmd/pt/model/descriptor/sezm_nn/so2.py b/deepmd/pt/model/descriptor/sezm_nn/so2.py index 0a3418efcf..4803b5d7c5 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/so2.py +++ b/deepmd/pt/model/descriptor/sezm_nn/so2.py @@ -1738,8 +1738,8 @@ def __init__( # scatter/gather softmax chain of the attention weights, sharing the # destination-sorted view with the flash aggregation; its backward and # hand-derived second order keep the force-loss trace from expanding - # the chain into materialized surfaces and serialized scatters. The - # source-gated (SFPG) form and fp64 compute keep the reference path. + # the chain into materialized surfaces and serialized scatters. fp64 + # compute keeps the reference path. self._segment_softmax_fn = None if ( max(self.triton_infer_level, self.triton_train_level) >= 1 @@ -1838,9 +1838,9 @@ def forward_envelope( Reduce the edge messages with the scalar envelope weight. The attention-free path: an envelope-weighted scatter add followed by the - degree normalization. Folding the Source Freeze Propagation Gate into the - envelope keeps the operation count unchanged; ``edge_src_gate`` is ``None`` - outside bridging mode, where the branch disappears. + degree normalization. In bridging mode the envelope already carries the + Source Freeze Propagation Gate, so a muted source drops out of both the + sum and the normalization. Parameters ---------- @@ -1860,12 +1860,8 @@ def forward_envelope( x_message, _ = self.edge_message(x, edge_cache, radial_feat) # (E, D, C_wide) - # === Step 2. Envelope weighting, with the source gate folded in === - edge_weight = edge_cache.edge_env # (E, 1) - edge_src_gate = edge_cache.edge_src_gate - if edge_src_gate is not None: - edge_weight = edge_weight * edge_src_gate.to(dtype=edge_weight.dtype) - x_message = x_message * edge_weight.unsqueeze(-1) # (E, D, C_wide) + # === Step 2. Envelope weighting === + x_message = x_message * edge_cache.edge_env.unsqueeze(-1) # (E, D, C_wide) # === Step 3. Destination reduction and degree normalization === out = x.new_zeros(x.shape, dtype=self.compute_dtype) @@ -1909,14 +1905,7 @@ def forward_attention( ) # (N, Fa, Ca) # === Step 2. Backend dispatch === - # The fused CUDA operator computes the attention weights itself, so it - # does not serve the bridging mode, whose source gate reshapes the - # softmax normalization. - run_cuda = ( - self._cuda_conv_fn is not None - and not self.training - and edge_cache.edge_src_gate is None - ) + run_cuda = self._cuda_conv_fn is not None and not self.training run_flash = ( self._flash_atten_fn is not None and (self._flash_atten_trains or not self.training) @@ -2166,12 +2155,12 @@ def attention_weights( """ Build envelope-gated attention weights from the scalar channels. - The softmax takes ``src_weight`` so that the Source Freeze Propagation - Gate enters both the numerator and the denominator. A muted source - (``eta_src = 0``) then drops out of the destination's normalization - entirely, which the frozen-zone invariance requires: post-multiplying the - weights alone would still leak the muted source through the shared - denominator. + The envelope enters both the numerator and the denominator of the + softmax. In bridging mode it carries the Source Freeze Propagation + Gate, so a muted source (``eta_src = 0``) drops out of the + destination's normalization entirely, which the frozen-zone invariance + requires: post-multiplying the weights alone would still leak the muted + source through the shared denominator. Parameters ---------- @@ -2212,11 +2201,9 @@ def attention_weights( ) # (E, F, H) # === Step 3. Envelope-gated segment softmax with a null mass === - edge_src_gate = edge_cache.edge_src_gate n_nodes = x_l0_node.shape[0] if ( self._segment_softmax_fn is not None - and edge_src_gate is None and attn_logits.is_cuda and active_triton_level(self) >= 1 ): @@ -2251,9 +2238,6 @@ def attention_weights( n_nodes=n_nodes, z_bias_raw=self.adamw_attn_z_bias_raw, eps=self.eps, - src_weight=( - None if edge_src_gate is None else edge_src_gate.to(dtype=compute_dtype) - ), ) # (E, F, H) def attention_head_gate(self, x_l0_node: torch.Tensor) -> torch.Tensor: diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 27c2d05af4..d2e0a09e96 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -37,11 +37,16 @@ preprocess_shared_params, ) from deepmd.utils.bridging import ( + DEFAULT_FRACTION_INNER, + DEFAULT_FRACTION_OUTER, + concise_bridging_window, expand_bridging_method, + resolve_bridging_window, route_canonical_learned_options, ) from deepmd.utils.spin import ( Spin, + env_protection_descriptors, ) from .dipole_model import ( @@ -185,11 +190,9 @@ def get_spin_model(model_params: dict) -> SpinModel: exclude_types=model_params.get("atom_exclude_types", None) ) model_params["atom_exclude_types"] = atom_exclude_types - if ( - "env_protection" not in model_params["descriptor"] - or model_params["descriptor"]["env_protection"] == 0.0 - ): - model_params["descriptor"]["env_protection"] = 0.01 + for descriptor in env_protection_descriptors(model_params["descriptor"]): + if "env_protection" not in descriptor or descriptor["env_protection"] == 0.0: + descriptor["env_protection"] = 0.01 if model_params["descriptor"]["type"] in ["se_e2_a"]: # only expand sel for se_e2_a model_params["descriptor"]["sel"] += model_params["descriptor"]["sel"] @@ -355,8 +358,14 @@ def _get_bridged_linear_model(model_params: dict) -> BaseModel: learned["atom_exclude_types"] = model_params.get("atom_exclude_types", []) learned["pair_exclude_types"] = model_params.get("pair_exclude_types", []) learned["bridging_method"] = inner_cfg.get("mode", "zbl") - learned["bridging_r_inner"] = float(inner_cfg.get("r_inner", 0.5)) - learned["bridging_r_outer"] = float(inner_cfg.get("r_outer", 0.8)) + learned["bridging_fraction_inner"] = inner_cfg.get( + "fraction_inner", DEFAULT_FRACTION_INNER + ) + learned["bridging_fraction_outer"] = inner_cfg.get( + "fraction_outer", DEFAULT_FRACTION_OUTER + ) + learned["bridging_r_inner"] = inner_cfg.get("r_inner") + learned["bridging_r_outer"] = inner_cfg.get("r_outer") if "spin" in model_params: learned["spin"] = model_params["spin"] return get_sezm_spin_model(learned) @@ -446,6 +455,24 @@ def get_standard_model(model_params: dict) -> BaseModel: return model +def _bridging_window_options(model_params: dict) -> dict: + """ + Resolve the bridging window of a SeZM model section. + + Parameters + ---------- + model_params : dict + A DPA4/SeZM model section carrying the concise bridging keys. + + Returns + ------- + dict + The window options of the learned descriptor, as + :func:`deepmd.utils.bridging.resolve_bridging_window` spells them. + """ + return resolve_bridging_window(concise_bridging_window(model_params)) + + def _reconcile_sezm_pair_exclude_types(model_params: dict) -> list[list[int]]: """Reconcile ``pair_exclude_types`` with ``descriptor.exclude_types``. @@ -504,12 +531,10 @@ def get_sezm_model(model_params: dict) -> BaseModel: # === Bridging parameters === bridging_method = str(model_params.get("bridging_method", "none")).upper() - bridging_r_inner = float(model_params.get("bridging_r_inner", 0.5)) - bridging_r_outer = float(model_params.get("bridging_r_outer", 0.8)) + bridging_window = _bridging_window_options(model_params) # Only inject bridging parameters when bridging is enabled. if bridging_method != "NONE": - model_params["descriptor"]["inner_clamp_r_inner"] = bridging_r_inner - model_params["descriptor"]["inner_clamp_r_outer"] = bridging_r_outer + model_params["descriptor"].update(bridging_window) descriptor = BaseDescriptor(**model_params["descriptor"]) @@ -562,8 +587,6 @@ def get_sezm_model(model_params: dict) -> BaseModel: use_compile=use_compile, enable_tf32=enable_tf32, bridging_method=bridging_method, - bridging_r_inner=bridging_r_inner, - bridging_r_outer=bridging_r_outer, ) model.model_def_script = json.dumps(model_params_old) return model @@ -620,11 +643,9 @@ def _get_sezm_native_spin_model(model_params: dict) -> BaseModel: # === Bridging parameters (no virtual atoms, so ZBL needs no masking) === bridging_method = str(model_params.get("bridging_method", "none")).upper() - bridging_r_inner = float(model_params.get("bridging_r_inner", 0.5)) - bridging_r_outer = float(model_params.get("bridging_r_outer", 0.8)) + bridging_window = _bridging_window_options(model_params) if bridging_method != "NONE": - model_params["descriptor"]["inner_clamp_r_inner"] = bridging_r_inner - model_params["descriptor"]["inner_clamp_r_outer"] = bridging_r_outer + model_params["descriptor"].update(bridging_window) descriptor = BaseDescriptor(**model_params["descriptor"]) @@ -658,8 +679,6 @@ def _get_sezm_native_spin_model(model_params: dict) -> BaseModel: use_compile=use_compile, enable_tf32=enable_tf32, bridging_method=bridging_method, - bridging_r_inner=bridging_r_inner, - bridging_r_outer=bridging_r_outer, spin=spin, ) model.model_def_script = json.dumps(model_params_old) @@ -699,14 +718,6 @@ def _get_sezm_virtual_spin_model(model_params: dict) -> BaseModel: model_params["descriptor"]["ntypes"] = ntypes model_params["descriptor"]["type_map"] = copy.deepcopy(model_params["type_map"]) - # === Bridging parameters === - bridging_method = str(model_params.get("bridging_method", "none")).upper() - bridging_r_inner = float(model_params.get("bridging_r_inner", 0.5)) - bridging_r_outer = float(model_params.get("bridging_r_outer", 0.8)) - if bridging_method != "NONE": - model_params["descriptor"]["inner_clamp_r_inner"] = bridging_r_inner - model_params["descriptor"]["inner_clamp_r_outer"] = bridging_r_outer - descriptor = BaseDescriptor(**model_params["descriptor"]) fitting_net = copy.deepcopy(model_params["fitting_net"]) @@ -737,9 +748,7 @@ def _get_sezm_virtual_spin_model(model_params: dict) -> BaseModel: data_stat_protect=data_stat_protect, use_compile=use_compile, enable_tf32=enable_tf32, - bridging_method=bridging_method, - bridging_r_inner=bridging_r_inner, - bridging_r_outer=bridging_r_outer, + bridging_method=str(model_params.get("bridging_method", "none")), real_sel=real_sel_list, spin=spin, ) diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py index 264c01e8ee..5048fbf8ab 100644 --- a/deepmd/pt/model/model/dp_model.py +++ b/deepmd/pt/model/model/dp_model.py @@ -1,5 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) + from deepmd.pt.model.descriptor.base_descriptor import ( BaseDescriptor, ) @@ -50,3 +54,15 @@ def get_fitting_net(self): # noqa: ANN201 def get_descriptor(self): # noqa: ANN201 """Get the descriptor.""" return self.atomic_model.descriptor + + def compute_fitting_input_stats( + self, sampled_func: Callable[[], list[dict]] + ) -> None: + """Recompute the input statistics of the fitting network from a sample. + + Parameters + ---------- + sampled_func + The lazy sampled function to get data frames from different data systems. + """ + self.get_fitting_net().compute_input_stats(sampled_func) diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index 91ff826a3e..b1fdf148db 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -472,6 +472,9 @@ from torch import Tensor from deepmd.dpmodel.utils.neighbor_list import EdgeNeighborList, NeighborList +from deepmd.dpmodel.atomic_model.inner_potential import ( + InnerPotential as InnerPotentialDP, +) from deepmd.pt.model.atomic_model.sezm_atomic_model import ( SeZMAtomicModel, ) @@ -652,10 +655,11 @@ def _sezm_structure_key(model: SeZMModel) -> tuple[Any, ...]: bool(descriptor.random_gamma), descriptor.charge_spin_embedding is not None, descriptor.spin_embedding is not None, - descriptor.inner_clamp is not None, + descriptor.bridging_clamp is not None, descriptor.bridging_switch is not None, - descriptor.inner_clamp_r_inner, - descriptor.inner_clamp_r_outer, + descriptor.bridging_f_inner, + descriptor.bridging_f_outer, + descriptor.bridging_scale, int(descriptor.get_dim_chg_spin()), ) fitting_state = ( @@ -666,8 +670,6 @@ def _sezm_structure_key(model: SeZMModel) -> tuple[Any, ...]: model_state = ( str(model.bridging_method), model.inter_potential is not None, - float(model.bridging_r_inner), - float(model.bridging_r_outer), tuple(model.get_type_map()), ) return ( @@ -709,8 +711,6 @@ def __init__( use_compile: bool = False, enable_tf32: bool = True, bridging_method: str = "none", - bridging_r_inner: float = 0.5, - bridging_r_outer: float = 0.8, lora: dict[str, Any] | None = None, **kwargs: Any, ) -> None: @@ -766,8 +766,6 @@ def __init__( # === Bridging (optional short-range zone bridging) === self.bridging_method: str = str(bridging_method).upper() - self.bridging_r_inner = float(bridging_r_inner) - self.bridging_r_outer = float(bridging_r_outer) self.inter_potential: InnerPotential | None = ( InnerPotential(type_map=self.get_type_map(), mode=self.bridging_method) if self.bridging_method != "NONE" @@ -1560,7 +1558,7 @@ def core_compute( # without ``comm_dict``, and ``extended_coord`` only supplies the device. fitting_net = self.atomic_model.fitting_net with nvtx_range("SeZM/descriptor"): - descriptor, _, vacuum = descriptor_model.forward_with_edges( + descriptor, _, vacuum, node_gate = descriptor_model.forward_with_edges( extended_coord=coord, extended_atype=descriptor_atype, edge_index=edge_index, @@ -1586,6 +1584,7 @@ def core_compute( aparam=aparam, vacuum_descriptor=vacuum, return_atomic_feature=embedding_only, + node_gate=node_gate, ) # === Embedding short circuit === @@ -1625,7 +1624,6 @@ def core_compute( atype_flat=descriptor_atype.reshape(-1), edge_mask=inter_potential_edge_mask, n_node=nf * nloc, - real_type_count=self._get_inter_potential_real_type_count(), ).view(nf, nloc, 1) # === Step 5. Apply atom mask to the complete physical output === @@ -1764,7 +1762,7 @@ def core_compute_dens( # === Step 3. Descriptor forward with force embedding === with nvtx_range("SeZM/descriptor_dens"): - descriptor, latent, vacuum = descriptor_model.forward_with_edges( + descriptor, latent, vacuum, node_gate = descriptor_model.forward_with_edges( extended_coord=extended_coord[:, :nloc, :], extended_atype=atype, edge_index=edge_index, @@ -1788,6 +1786,7 @@ def core_compute_dens( aparam=aparam, vacuum_descriptor=vacuum, return_components=True, + node_gate=node_gate, ) return torch.cat( [ @@ -3277,6 +3276,12 @@ def _make_inter_potential_edge_mask( ) -> torch.Tensor: """Build the complete validity mask for analytical pair energies. + The analytical term receives the pair exclusions of the descriptor and + skips every edge that touches a phantom atom. Atom exclusions do not + enter: they act on atomic outputs, so the final atom mask drops the + half-energy of an excluded atom while its partner keeps its own half, + as the composition of the other backends does. + Parameters ---------- atype @@ -3298,7 +3303,6 @@ def _make_inter_potential_edge_mask( src = edge_index[0] dst = edge_index[1] - atype_flat = atype.reshape(-1) real_atom_flat = real_atom.reshape(-1) keep = ( edge_mask @@ -3308,23 +3312,9 @@ def _make_inter_potential_edge_mask( descriptor = self.atomic_model.descriptor if descriptor.exclude_types: - keep = keep & descriptor._edge_type_keep_mask(atype_flat, src, dst) - - atom_excl = self.atomic_model.atom_excl - if atom_excl is not None: - atom_is_included = atom_excl(atype).to(torch.bool).reshape(-1) - keep = ( - keep - & atom_is_included.index_select(0, src) - & atom_is_included.index_select(0, dst) - ) - + keep = keep & descriptor._edge_type_keep_mask(atype.reshape(-1), src, dst) return keep - def _get_inter_potential_real_type_count(self) -> int: - """Return the real-type count used to mask analytical pair potentials.""" - return len(self.get_type_map()) - # ========================================================================= # Type and Output Metadata # ========================================================================= @@ -3402,8 +3392,6 @@ def serialize(self) -> dict[str, Any]: "type": self.model_type, "atomic_model": self.atomic_model.serialize(), "bridging_method": self.bridging_method, - "bridging_r_inner": self.bridging_r_inner, - "bridging_r_outer": self.bridging_r_outer, "lora": self.lora_config, } @@ -3486,47 +3474,26 @@ def tf32_precision_ctx(self) -> Generator[None, None, None]: # InnerPotential: analytical pair potentials for bridging # ============================================================================= -# 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) +class InnerPotential(torch.nn.Module): + r""" + Analytical pair potential module for Zone bridging. -# Physical constants -_KE_EV_A = 14.3996 # Coulomb constant in eV·Å -_A_BOHR = 0.5291772109 # Bohr radius in Å + Every supported potential is a screened Coulomb series in the pair + separation, + .. math:: -class InnerPotential(torch.nn.Module): - """ - Analytical pair potential module for Zone bridging. + V_{ab}(r) = \frac1r \sum_{k=1}^{4} A_{abk} e^{-c_{abk} r}, - Supports the Ziegler-Biersack-Littmark (ZBL) screened nuclear repulsion - potential. Designed to be extensible to other analytical forms (LJ, Morse, - etc.) through the ``mode`` parameter. + so one table of eight constants per ordered type pair describes all of + them. ``"ZBL"`` builds that table from the Ziegler-Biersack-Littmark + universal screening function, and ``"NLH"`` from the bundled + Nordlund-Lehtola-Hobler coefficients; the table itself is built by the + backend-agnostic :class:`~deepmd.dpmodel.atomic_model.inner_potential.InnerPotential`, + so both backends evaluate the same constants. - Each pair (i, j) contributes ``V_ZBL(r_ij) / 2`` to both atom i and atom j, + Each pair (i, j) contributes ``V(r_ij) / 2`` to both atom i and atom j, avoiding double-counting from the symmetric neighbor list. Parameters @@ -3535,7 +3502,7 @@ class InnerPotential(torch.nn.Module): 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. + Potential formula, ``"zbl"`` or ``"nlh"``. Case-insensitive. Raises ------ @@ -3547,50 +3514,23 @@ class InnerPotential(torch.nn.Module): def __init__(self, type_map: list[str], mode: str = "zbl") -> None: super().__init__() mode = mode.upper() - if mode != "ZBL": + if mode not in InnerPotentialDP.SUPPORTED_MODES: raise ValueError(f"Unknown InnerPotential mode: {mode}") self.mode = mode - self.ntypes_real = len(type_map) - - 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) + self.ntypes = len(type_map) + + # Derived from ``type_map`` and ``mode``, hence rebuilt on construction + # and kept out of the checkpoint. self.register_buffer( - "atomic_numbers", - torch.tensor(atomic_numbers, dtype=torch.float64, device=env.DEVICE), + "series_table", + torch.tensor( + InnerPotentialDP.series_table_from_type_map(type_map, mode=mode), + dtype=torch.float64, + device=env.DEVICE, + ), + persistent=False, ) - def _zbl_pair_energy( - self, - r: torch.Tensor, - zi: torch.Tensor, - zj: torch.Tensor, - ) -> torch.Tensor: - """ - Compute ZBL pair energy for given distances and nuclear charges. - - Parameters - ---------- - r : torch.Tensor - Pair distances with shape (...) in Å. - zi : torch.Tensor - Nuclear charge of atom i with shape (...). - zj : torch.Tensor - Nuclear charge of atom j with shape (...). - - Returns - ------- - torch.Tensor - Pair energies with shape (...) in eV. - """ - a_screen = 0.88534 * _A_BOHR / (zi.pow(0.23) + zj.pow(0.23)) - x = r / a_screen - phi = sum(a * torch.exp(-b * x) for a, b in zip(_ZBL_A_COEFF, _ZBL_B_COEFF)) - return _KE_EV_A * zi * zj / r * phi - def forward( self, edge_vec: torch.Tensor, @@ -3598,7 +3538,6 @@ def forward( atype_flat: torch.Tensor, edge_mask: torch.Tensor, n_node: int, - real_type_count: int | None = None, ) -> torch.Tensor: """ Compute per-atom pair energy from the sparse edge list. @@ -3620,42 +3559,25 @@ def forward( Boolean mask with shape (E,). True means valid edge. n_node : int Number of flattened local nodes. - real_type_count - Number of real atom types. Types ``>= real_type_count`` are - virtual spin types: edges touching them are masked out. If - ``None``, all configured types are real. Returns ------- torch.Tensor Per-atom pair energy with shape (1, N, 1) in eV. """ - if real_type_count is None: - real_type_count = self.ntypes_real + stride = self.ntypes + 1 src = edge_index[0].to(dtype=torch.long) dst = edge_index[1].to(dtype=torch.long) - - # Wrap virtual spin types (>= real_type_count) back onto their real - # parent type so the Z lookup never indexes out of range; their edges - # are masked out below regardless. - atype_for_z = atype_flat.clamp(min=0) - atype_for_z = torch.where( - atype_for_z >= real_type_count, - atype_for_z - real_type_count, - atype_for_z, - ) + # Padded atoms carry type -1; their edges are masked below, so the + # table lookup only needs a valid index for them. + atype_row = atype_flat.clamp(min=0) r = edge_vec.to(dtype=torch.float64).norm(dim=-1).clamp(min=1e-10) # (E,) - z_all = self.atomic_numbers[atype_for_z] # (N,) - zi = z_all[src] # (E,) - zj = z_all[dst] # (E,) - - pair_e = self._zbl_pair_energy(r, zi, zj) # (E,) - # Drop padded edges and any edge touching a virtual (spin) node. - node_is_real = atype_flat < real_type_count # (N,) - edge_is_real = node_is_real[src] & node_is_real[dst] # (E,) - valid = edge_mask & edge_is_real - pair_e = pair_e * valid.to(dtype=pair_e.dtype) + # Row ``center * (T + 1) + neighbor``, the ordered-pair index the fused + # kernels use, with the destination as the center. + row = self.series_table[atype_row[dst] * stride + atype_row[src]] # (E, 8) + pair_e = (row[:, :4] * torch.exp(-row[:, 4:] * r[:, None])).sum(-1) / r # (E,) + pair_e = pair_e * edge_mask.to(dtype=pair_e.dtype) # Half contribution to each destination atom atom_energy = torch.zeros(n_node, dtype=pair_e.dtype, device=pair_e.device) diff --git a/deepmd/pt/model/model/sezm_spin_model.py b/deepmd/pt/model/model/sezm_spin_model.py index ace082b02b..86b7eeca3d 100644 --- a/deepmd/pt/model/model/sezm_spin_model.py +++ b/deepmd/pt/model/model/sezm_spin_model.py @@ -27,7 +27,6 @@ BaseModel, ) from deepmd.pt.model.model.sezm_model import ( - InnerPotential, SeZMModel, ) from deepmd.pt.model.model.spin_model import ( @@ -76,10 +75,16 @@ def __init__( real_sel: list[int], **kwargs: Any, ) -> None: - # Delay InnerPotential construction until ntypes_real is available. - bridging_method = str(kwargs.pop("bridging_method", "none")).upper() - kwargs["bridging_method"] = "none" - + # A magnetic atom lies inside the bridging window of its own virtual + # partner, so the window mechanics would freeze every magnetic atom; + # the native spin scheme carries no virtual atoms and is the bridged + # spin model. + if str(kwargs.get("bridging_method", "none")).upper() != "NONE": + raise ValueError( + "Analytical bridging is not supported by the deepspin " + "(virtual-atom) scheme; use spin scheme 'native' for a bridged " + "DPA4/SeZM spin model." + ) super().__init__(*args, **kwargs) self.spin = spin self.ntypes_real = self.spin.ntypes_real @@ -95,13 +100,6 @@ def __init__( persistent=False, ) - self.bridging_method = bridging_method - self.inter_potential = ( - InnerPotential(type_map=self.get_type_map(), mode=self.bridging_method) - if self.bridging_method != "NONE" - else None - ) - # ========================================================================= # Forward Methods # ========================================================================= @@ -601,10 +599,6 @@ def format_nlist( extra_nlist_sort=extra_nlist_sort, ) - def _get_inter_potential_real_type_count(self) -> int: - """Return the number of real types for real-only ZBL masking.""" - return self.ntypes_real - def _get_spin_sampled_func( self, sampled_func: Callable[[], list[dict[str, Any]]] ) -> Callable[[], list[dict[str, Any]]]: diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index c59a336d71..762a2f0cf2 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -481,8 +481,9 @@ def __init__( if not self.mixed_types: assert self.ntypes == bias_atom_e.shape[0], "Element count mismatches!" self.register_buffer("bias_atom_e", bias_atom_e) - # A deployment constant; see :meth:`fold_vacuum_reference`. + # Deployment constants; see :meth:`fold_vacuum_reference`. self.register_buffer("vacuum_table", None, persistent=False) + self.register_buffer("folded_reference", None, persistent=False) if self.numb_fparam > 0: self.register_buffer( @@ -594,6 +595,7 @@ def change_type_map( self.bias_atom_e = self.bias_atom_e[remap_index] # the stored references belong to the old type map self.vacuum_table = None + self.folded_reference = None def serialize(self) -> dict: """Serialize the fitting to dict.""" @@ -943,12 +945,14 @@ def fold_vacuum_reference(self, vacuum_descriptor: torch.Tensor) -> None: Under uniform conditioning the reference output of an atom is a constant of its type, so subtracting it from ``bias_atom_e`` yields the same outputs as the referenced forward and the option is switched - off. With frame or atomic parameters the reference output varies - between atoms, so the vacuum descriptor is stored instead and the - forward evaluates the references from the stored table. The table is - a deployment constant: an exported model bakes it, checkpoints leave - it out, and a fitting loaded from a checkpoint takes the reference - from the descriptor again. + off; the subtracted outputs are kept as ``folded_reference`` so that + :meth:`readout_reference` still knows the bias the fold moved. With + frame or atomic parameters the reference output varies between atoms, + so the vacuum descriptor is stored instead and the forward evaluates + the references from the stored table. Both are deployment constants: + an exported model bakes them, checkpoints leave them out, and a + fitting loaded from a checkpoint takes the reference from the + descriptor again. Parameters ---------- @@ -960,14 +964,33 @@ def fold_vacuum_reference(self, vacuum_descriptor: torch.Tensor) -> None: return with torch.no_grad(): if self.uniform_conditioning(): - reference = self.vacuum_property(vacuum_descriptor) - self.bias_atom_e = self.bias_atom_e - reference.to( + reference = self.vacuum_property(vacuum_descriptor).to( self.bias_atom_e.dtype ) + self.folded_reference = reference.detach().clone() + self.bias_atom_e = self.bias_atom_e - reference self.vacuum_ref = False else: self.vacuum_table = vacuum_descriptor.detach().to(self.prec).clone() + def readout_reference(self) -> torch.Tensor: + """Per-type constant the readout gate of a bridging window fades to. + + It is the fitting bias ``bias_atom_e`` of every type, which under the + vacuum reference is the output of an isolated atom of that type. A + fold under uniform conditioning moves the reference outputs out of the + bias; the stored copy restores the bias the gate reads, so the gate + is the same function of the atom before and after the fold. + + Returns + ------- + torch.Tensor + The reference of every type with shape (ntypes, dim_out). + """ + if self.folded_reference is None: + return self.bias_atom_e + return self.bias_atom_e + self.folded_reference + def _forward_common( self, descriptor: torch.Tensor, @@ -979,6 +1002,7 @@ def _forward_common( aparam: torch.Tensor | None = None, vacuum_descriptor: torch.Tensor | None = None, return_atomic_feature: bool = False, + node_gate: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: # cast the input to internal precision xx = descriptor.to(self.prec) @@ -1078,6 +1102,16 @@ def _forward_common( outs = ( outs + atom_property ) # Shape is [nframes, natoms[0], net_dim_out] + # === Step 3. Readout gate of a bridging window === + # The learned part of the output is its deviation from the bias of + # the type (see :meth:`readout_reference`), which is the output of an + # isolated atom whenever the fitting references its atoms. A bridged + # descriptor hands over the per-atom source gate, and the learned part + # fades with it: an atom of a frozen pair keeps only the bias of its + # type. + if node_gate is not None: + reference = self.readout_reference()[atype].to(self.prec) + outs = reference + node_gate.to(self.prec) * (outs - reference) # nf x nloc mask = self.emask(atype).to(torch.bool) # nf x nloc x nod diff --git a/deepmd/pt/model/task/invar_fitting.py b/deepmd/pt/model/task/invar_fitting.py index 8137772a3a..0fde6071a1 100644 --- a/deepmd/pt/model/task/invar_fitting.py +++ b/deepmd/pt/model/task/invar_fitting.py @@ -192,6 +192,7 @@ def forward( aparam: torch.Tensor | None = None, vacuum_descriptor: torch.Tensor | None = None, return_atomic_feature: bool = False, + node_gate: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Based on embedding net output, alculate total energy. @@ -202,6 +203,8 @@ def forward( shape [ntypes, self.dim_descrpt], required by ``vacuum_ref``. - return_atomic_feature: also return the last hidden activation under the ``atomic_feature`` key. + - node_gate: per-atom source gate of a bridged descriptor with shape + [nframes, natoms[0], 1]; the learned part of the output fades with it. Returns ------- @@ -217,6 +220,7 @@ def forward( aparam, vacuum_descriptor=vacuum_descriptor, return_atomic_feature=return_atomic_feature, + node_gate=node_gate, ) result = {self.var_name: out[self.var_name].to(env.GLOBAL_PT_FLOAT_PRECISION)} if return_atomic_feature: diff --git a/deepmd/pt/model/task/sezm_ener.py b/deepmd/pt/model/task/sezm_ener.py index 2358c1942e..67633514eb 100644 --- a/deepmd/pt/model/task/sezm_ener.py +++ b/deepmd/pt/model/task/sezm_ener.py @@ -647,6 +647,7 @@ def _forward_common( aparam: torch.Tensor | None = None, vacuum_descriptor: torch.Tensor | None = None, return_atomic_feature: bool = False, + node_gate: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Run the SeZM fitting path with optional case FiLM.""" if not self.case_film_embd: @@ -660,6 +661,7 @@ def _forward_common( aparam, vacuum_descriptor=vacuum_descriptor, return_atomic_feature=return_atomic_feature, + node_gate=node_gate, ) return self._forward_case_film( descriptor, @@ -668,6 +670,7 @@ def _forward_common( aparam, vacuum_descriptor=vacuum_descriptor, return_atomic_feature=return_atomic_feature, + node_gate=node_gate, ) def _forward_case_film( @@ -678,6 +681,7 @@ def _forward_case_film( aparam: torch.Tensor | None = None, vacuum_descriptor: torch.Tensor | None = None, return_atomic_feature: bool = False, + node_gate: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """ Forward path for SeZM case FiLM. @@ -698,6 +702,9 @@ def _forward_case_film( return_atomic_feature When True, also return the last hidden activation under the ``atomic_feature`` key. + node_gate + Per-atom source gate of a bridged descriptor with shape + (nf, nloc, 1); the learned part of the output fades with it. Returns ------- @@ -739,6 +746,11 @@ def _forward_case_film( fitting(xx_vac, self.case_embd), atype ) outs = atom_property + self.bias_atom_e[atype].to(self.prec) + # The readout gate of a bridging window fades the learned part of the + # output, its deviation from the bias of the type. + if node_gate is not None: + reference = self.readout_reference()[atype].to(self.prec) + outs = reference + node_gate.to(self.prec) * (outs - reference) mask = self.emask(atype).to(torch.bool) outs = torch.where(mask[:, :, None], outs, 0.0) diff --git a/deepmd/pt/nvalchemi/dpa4wrapper.py b/deepmd/pt/nvalchemi/dpa4wrapper.py index 720d1b5472..a49302ffb6 100644 --- a/deepmd/pt/nvalchemi/dpa4wrapper.py +++ b/deepmd/pt/nvalchemi/dpa4wrapper.py @@ -58,8 +58,10 @@ nn, ) -from deepmd.pt.model.model.sezm_model import ( +from deepmd.dpmodel.atomic_model.inner_potential import ( ELEMENT_TO_Z, +) +from deepmd.pt.model.model.sezm_model import ( SeZMModel, ) diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 7a7210962e..7efc76e4d9 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -103,7 +103,7 @@ ) from deepmd.pt.utils.stat import ( make_stat_input, - min_pair_dist_frame_mask, + pair_margin_frame_mask, scan_redu_stats, select_batch_frames, ) @@ -132,6 +132,7 @@ from deepmd.utils.data import ( DataRequirementItem, has_data_requirement, + min_pair_dist_requirement, ) from deepmd.utils.finetune import ( warn_configuration_mismatch_during_finetune, @@ -519,25 +520,22 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: ) # Data + # Radius of the pair-clearance filter, per task. + self.min_pair_dist_by_task: dict[str, float] = {} if not self.multi_task: # add data requirement for labels data_requirement = self.loss.label_requirement data_requirement += get_additional_data_requirement(self.model) - min_pair_dist = float( - training_params.get("training_data", {}).get("min_pair_dist", 0.0) + pair_filter = min_pair_dist_requirement( + model_params, + float( + training_params.get("training_data", {}).get("min_pair_dist", 0.0) + ), ) - if min_pair_dist > 0.0: - data_requirement.append( - DataRequirementItem( - "min_pair_dist", - ndof=1, - atomic=False, - must=False, - high_prec=False, - default=min_pair_dist, - source_policy="derived", - ) - ) + min_pair_dist = 0.0 if pair_filter is None else float(pair_filter.default) + self.min_pair_dist_by_task["Default"] = min_pair_dist + if pair_filter is not None: + data_requirement.append(pair_filter) training_data.add_data_requirement(data_requirement) if validation_data is not None: validation_data.add_data_requirement( @@ -601,23 +599,20 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: data_requirement += get_additional_data_requirement( self.model[model_key] ) - min_pair_dist = float( - training_params["data_dict"][model_key] - .get("training_data", {}) - .get("min_pair_dist", 0.0) + pair_filter = min_pair_dist_requirement( + model_params["model_dict"][model_key], + float( + training_params["data_dict"][model_key] + .get("training_data", {}) + .get("min_pair_dist", 0.0) + ), ) - if min_pair_dist > 0.0: - data_requirement.append( - DataRequirementItem( - "min_pair_dist", - ndof=1, - atomic=False, - must=False, - high_prec=False, - default=min_pair_dist, - source_policy="derived", - ) - ) + min_pair_dist = ( + 0.0 if pair_filter is None else float(pair_filter.default) + ) + self.min_pair_dist_by_task[model_key] = min_pair_dist + if pair_filter is not None: + data_requirement.append(pair_filter) training_data[model_key].add_data_requirement(data_requirement) if validation_data[model_key] is not None: validation_data[model_key].add_data_requirement( @@ -738,24 +733,12 @@ def epoch_length(model_key: str) -> int: self.nonfinite_grad_guard = NonFiniteGradGuard() self.lr_schedule = get_lr(config["learning_rate"]) - # Minimum pairwise distance for filtering unphysical frames during training. - if self.multi_task: - self.min_pair_dist: float | dict[str, float] = { - model_key: float( - training_params["data_dict"][model_key] - .get("training_data", {}) - .get("min_pair_dist", 0.0) - ) - for model_key in self.model_keys - } - else: - self.min_pair_dist = float( - training_params.get("training_data", {}).get("min_pair_dist", 0.0) - ) - self.has_min_pair_filter = ( - any(value > 0.0 for value in self.min_pair_dist.values()) - if isinstance(self.min_pair_dist, dict) - else self.min_pair_dist > 0.0 + # Radius the pair-clearance filter was registered with. The + # filter reads it as an enable flag, the frames carrying a margin that + # is already measured against that radius. It is resolved once with the + # data requirement, so the two cannot disagree. + self.has_min_pair_filter = any( + value > 0.0 for value in self.min_pair_dist_by_task.values() ) if self.has_min_pair_filter: if self.multi_task: @@ -2161,12 +2144,6 @@ def save_ema_model_merged( use_ema_weights=True, ) - def _get_min_pair_dist(self, task_key: str) -> float: - """Return the minimum pair distance configured for one task.""" - if isinstance(self.min_pair_dist, dict): - return self.min_pair_dist[task_key] - return self.min_pair_dist - def _next_training_batch( self, task_key: str, @@ -2207,8 +2184,12 @@ def _next_training_batch( raise RuntimeError( "Unable to collect a globally valid training batch for task " - f"{task_key!r} after {max_attempts} attempts with " - f"min_pair_dist={self._get_min_pair_dist(task_key)}." + f"{task_key!r} after {max_attempts} attempts: no frame keeps every " + "atom pair beyond the frame-filter radius. A bridged model filters at " + "the midpoint of its bridging window, so lower " + "`bridging_fraction_inner` or `bridging_fraction_outer` there, or " + "`training_data.min_pair_dist` on any other model, or clean the " + "dataset." ) def get_data( @@ -2224,10 +2205,8 @@ def get_data( return {}, {}, {} batch_data = next(iterator) # === Filter frames with atoms too close (training only) === - min_pair_dist = self._get_min_pair_dist(task_key) - valid_mask = ( - min_pair_dist_frame_mask(batch_data, min_pair_dist) if is_train else None - ) + filter_enabled = is_train and self.min_pair_dist_by_task[task_key] > 0.0 + valid_mask = pair_margin_frame_mask(batch_data, filter_enabled) local_has_valid = valid_mask is None or bool(torch.any(valid_mask)) if not local_has_valid: return {}, {}, {} @@ -2349,10 +2328,36 @@ def print_on_training( fout.flush() +#: The world a validity group was drawn from, paired with that group. The flag +#: is a host value and its answer gates host control flow, so reducing it over a +#: CPU backend keeps it off the accelerator and out of its execution stream, +#: where reading it back would drain the pipeline once per step. +_VALIDITY_GROUP: tuple[Any, Any] | None = None + + +def _validity_group() -> Any: + """Return the CPU process group used for the frame-validity agreement. + + The group belongs to the world it was drawn from, so it is retained under + that world and rebuilt when a process enters a new one. + """ + global _VALIDITY_GROUP + world = dist.group.WORLD + if _VALIDITY_GROUP is None or _VALIDITY_GROUP[0] is not world: + group = ( + world if dist.get_backend() == "gloo" else dist.new_group(backend="gloo") + ) + _VALIDITY_GROUP = (world, group) + return _VALIDITY_GROUP[1] + + def all_ranks_have_valid_frames(local_has_valid: bool) -> bool: """ Return whether every distributed rank has a valid training frame. + The agreement travels over the CPU process group of :func:`_validity_group`, + so the flag it reduces lives on the host whatever device the model trains on. + Parameters ---------- local_has_valid @@ -2364,11 +2369,9 @@ def all_ranks_have_valid_frames(local_has_valid: bool) -> bool: ``True`` only when every rank reports a valid frame. """ all_ranks_have_valid = torch.tensor( - int(local_has_valid), - dtype=torch.int32, - device=DEVICE, + int(local_has_valid), dtype=torch.int32, device="cpu" ) - dist.all_reduce(all_ranks_have_valid, op=dist.ReduceOp.MIN) + dist.all_reduce(all_ranks_have_valid, op=dist.ReduceOp.MIN, group=_validity_group()) return bool(all_ranks_have_valid.item()) diff --git a/deepmd/pt/utils/dataset.py b/deepmd/pt/utils/dataset.py index 20a76a0e87..d3a3583938 100644 --- a/deepmd/pt/utils/dataset.py +++ b/deepmd/pt/utils/dataset.py @@ -74,6 +74,7 @@ def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> N default=data_item["default"], dtype=data_item["dtype"], output_natoms_for_type_sel=data_item["output_natoms_for_type_sel"], + length_scale=data_item["length_scale"], ) def preload_and_modify_all_data_torch(self) -> None: diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 310079221d..49ffd995cf 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -65,39 +65,52 @@ "_restore_observed_type_from_file", "_save_observed_type_to_file", "collect_observed_types", - "min_pair_dist_frame_mask", "observed_types_from_counts", + "pair_margin_frame_mask", "scan_redu_stats", "select_batch_frames", ] -def min_pair_dist_frame_mask( +def pair_margin_frame_mask( batch: dict[str, Any], - min_pair_dist: float, + enabled: bool, ) -> torch.Tensor | None: """ - Return the valid-frame mask for a minimum pair-distance threshold. + Return the valid-frame mask of the pair-clearance filter. + + The derived field holds each frame's margin against the thresholds its atom + pairs carry, so a frame is kept when that margin reaches one, whatever + length scale sized them. The registered filter radius is already + spent on that margin and takes no part in this comparison. Parameters ---------- batch Data batch containing frame-aligned tensors. - min_pair_dist - Minimum allowed pair distance in Å. + enabled + Whether the pair-clearance filter is configured for this data. Returns ------- torch.Tensor or None Boolean mask with shape (nframes), or ``None`` when filtering is - disabled or the distance field is unavailable. + disabled. + + Raises + ------ + RuntimeError + If the batch carries a placeholder instead of a derived margin, which + would otherwise accept every frame silently. """ - if min_pair_dist <= 0.0 or "min_pair_dist" not in batch: + if not enabled: return None - distances = batch["min_pair_dist"] - if not isinstance(distances, torch.Tensor): - return None - return distances.reshape(-1) >= float(min_pair_dist) + if not batch.get("find_pair_margin", False): + raise RuntimeError( + "The pair-clearance filter is enabled, but the data source did not " + "derive the pair margin of the batch." + ) + return batch["pair_margin"].reshape(-1) >= 1.0 def select_batch_frames( @@ -155,8 +168,10 @@ def make_stat_input( nbatches Maximum number of valid batches collected from each system. min_pair_dist - Minimum allowed pair distance in Å. Frames below the threshold are - excluded before statistics are accumulated. + Radius the training-frame filter was registered with. When + positive, frames holding a pair inside its window are dropped before + the statistics are accumulated, so they see the frames the optimizer + sees. Returns ------- @@ -186,7 +201,7 @@ def make_stat_input( except StopIteration: break scanned_batches += 1 - frame_mask = min_pair_dist_frame_mask(stat_data, min_pair_dist) + frame_mask = pair_margin_frame_mask(stat_data, min_pair_dist > 0.0) if frame_mask is not None and not torch.any(frame_mask): continue if frame_mask is not None: @@ -216,9 +231,8 @@ def make_stat_input( if min_pair_dist > 0.0: log.info( "Skipping data system %d in statistics because no frame " - "satisfies min_pair_dist=%s.", + "keeps every atom pair beyond the filter radius.", system_index, - min_pair_dist, ) else: log.info( @@ -285,8 +299,8 @@ def scan_redu_stats( intensive Whether the fitting target is intensive. min_pair_dist - Minimum allowed pair distance in Angstrom. Frames below the threshold - are excluded. + Radius the training-frame filter was registered with. When + positive, frames holding a pair inside its window are excluded. Returns ------- @@ -307,7 +321,7 @@ def scan_redu_stats( with torch.device("cpu"): for dataloader in dataloaders: for batch in _full_pass_loader(dataloader): - frame_mask = min_pair_dist_frame_mask(batch, min_pair_dist) + frame_mask = pair_margin_frame_mask(batch, min_pair_dist > 0.0) if frame_mask is not None: if not torch.any(frame_mask): continue diff --git a/deepmd/pt_expt/common.py b/deepmd/pt_expt/common.py index 42080a9142..1c19f7b1b0 100644 --- a/deepmd/pt_expt/common.py +++ b/deepmd/pt_expt/common.py @@ -17,6 +17,8 @@ from collections.abc import ( Callable, + Iterable, + Iterator, ) from functools import ( wraps, @@ -193,8 +195,37 @@ def _auto_wrap_native_op(value: NativeOP) -> torch.nn.Module: return wrapped_cls.deserialize(value.serialize()) +class BufferList(torch.nn.Module): + """Hold tensors as persistent buffers behind the interface of a list. + + An index table moves with its module and enters the ``state_dict``, but no + optimizer updates it, which makes it a buffer rather than a parameter. The + entries are registered under their positions, the keys a + ``torch.nn.ParameterList`` uses. + + Parameters + ---------- + tensors : Iterable[torch.Tensor] + The entries, in order. + """ + + def __init__(self, tensors: Iterable[torch.Tensor]) -> None: + super().__init__() + for index, tensor in enumerate(tensors): + self.register_buffer(str(index), tensor) + + def __len__(self) -> int: + return len(self._buffers) + + def __getitem__(self, index: int) -> torch.Tensor: + return self._buffers[str(range(len(self))[index])] + + def __iter__(self) -> Iterator[torch.Tensor]: + return iter(self._buffers.values()) + + def _try_convert_list(name: str, value: list) -> torch.nn.Module | None: - """Try to convert a plain list to ModuleList or ParameterList. + """Try to convert a plain list to ModuleList, ParameterList or BufferList. Returns the converted container, or None if no conversion is needed (e.g., empty list, list of scalars/strings). @@ -222,19 +253,23 @@ def _try_convert_list(name: str, value: list) -> torch.nn.Module | None: ) converted.append(c) return torch.nn.ModuleList(converted) - # List of numpy arrays → ParameterList + # List of numpy arrays → ParameterList, or BufferList when every array is + # an index table: a module's parameters are what an optimizer receives, and + # sharded optimizers require them to share one floating-point type. if all(isinstance(v, np.ndarray) for v in value): from deepmd.pt_expt.utils import env # deferred - avoids circular import - params = [] - for v in value: - t = torch.as_tensor(v, device=env.DEVICE) - params.append( + tensors = [torch.as_tensor(v, device=env.DEVICE) for v in value] + if not any(t.is_floating_point() or t.is_complex() for t in tensors): + return BufferList(tensors) + return torch.nn.ParameterList( + [ torch.nn.Parameter( t, requires_grad=t.is_floating_point() or t.is_complex() ) - ) - return torch.nn.ParameterList(params) + for t in tensors + ] + ) return None @@ -344,7 +379,8 @@ class lists in ``CONFIG_DERIVED_ARRAYS`` becomes a NON-persistent buffer obj._buffers[name] = None return True, None - # list of modules / NativeOP / numpy arrays → ModuleList / ParameterList + # list of modules / NativeOP / numpy arrays → ModuleList / ParameterList / + # BufferList if isinstance(value, list) and "_modules" in obj.__dict__: converted_list = _try_convert_list(name, value) if converted_list is not None: diff --git a/deepmd/pt_expt/descriptor/dpa4.py b/deepmd/pt_expt/descriptor/dpa4.py index db843ffb01..546d503ae0 100644 --- a/deepmd/pt_expt/descriptor/dpa4.py +++ b/deepmd/pt_expt/descriptor/dpa4.py @@ -54,11 +54,11 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: return self.call(*args, **kwargs) -# InnerClamp/BridgingSwitch are parameter-free (scalar bridging radii only, +# InnerClamp/BridgingSwitch are parameter-free (the two window fractions only, # no serialize()); rebuild fresh from the stored constructor arguments. register_dpmodel_mapping( InnerClampDP, - lambda v: InnerClamp(v.r_inner, v.r_outer), + lambda v: InnerClamp(v.f_inner, v.f_outer), ) @@ -70,7 +70,7 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: register_dpmodel_mapping( BridgingSwitchDP, - lambda v: BridgingSwitch(v.r_inner, v.r_outer), + lambda v: BridgingSwitch(v.f_inner, v.f_outer), ) @@ -198,6 +198,11 @@ def _promote_trainable_tree(module: torch.nn.Module) -> torch.nn.Module: class DescrptDPA4(DescrptDPA4DP): _update_sel_cls = UpdateSel + #: The per-element length scales of the bridging window follow from + #: ``type_map``, which the constructor resolves on every load, so the + #: buffer they land in stays out of the state dict. + CONFIG_DERIVED_ARRAYS = ("contact_radius",) + def adam_route_patterns(self) -> list[str]: """ Name patterns, relative to the descriptor, of the tensors that take the @@ -214,16 +219,11 @@ def adam_route_patterns(self) -> list[str]: def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) # The fused convolution paths consume only the three structural rows of - # each Wigner degree block. Source-gated attention bypasses that fused - # convolution, so its dense per-edge rotations remain available. - self._wigner_free_conv = ( - self.bridging_switch is None - and bool(self.blocks) - and all( - getattr(block.so2_conv, "_cuda_conv_fn", None) is not None - and not block.so2_conv._cuda_conv_fn._compete - for block in self.blocks - ) + # each Wigner degree block. + self._wigner_free_conv = bool(self.blocks) and all( + getattr(block.so2_conv, "_cuda_conv_fn", None) is not None + and not block.so2_conv._cuda_conv_fn._compete + for block in self.blocks ) self._packed_wigner_train = bool(self.blocks) and all( getattr(block.so2_conv, "_cuda_value_train", None) is not None diff --git a/deepmd/pt_expt/descriptor/dpa4_nn/so2.py b/deepmd/pt_expt/descriptor/dpa4_nn/so2.py index 87e3344e7c..83948c6fbf 100644 --- a/deepmd/pt_expt/descriptor/dpa4_nn/so2.py +++ b/deepmd/pt_expt/descriptor/dpa4_nn/so2.py @@ -372,8 +372,8 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: # scatter/gather softmax chain of the attention weights, sharing the # destination-sorted view with the flash aggregation; its backward and # hand-derived second order keep the force-loss trace from expanding - # the chain into materialized surfaces and serialized scatters. The - # source-gated (SFPG) form and fp64 compute keep the reference path. + # the chain into materialized surfaces and serialized scatters. fp64 + # compute keeps the reference path. self._segment_softmax_fn = None if ( max(self.triton_infer_level, self.triton_train_level) >= 1 @@ -417,7 +417,6 @@ def _attention_softmax( active_level = _active_triton_level(self) if ( self._segment_softmax_fn is not None - and edge_cache.edge_src_gate is None and attn_logits.is_cuda and active_level >= 1 ): diff --git a/deepmd/pt_expt/descriptor/dpa4c.py b/deepmd/pt_expt/descriptor/dpa4c.py index 4fa1b17d66..a6482fbaa6 100644 --- a/deepmd/pt_expt/descriptor/dpa4c.py +++ b/deepmd/pt_expt/descriptor/dpa4c.py @@ -15,8 +15,10 @@ import torch +from deepmd.dpmodel.descriptor.dpa4_nn.radial import BridgingClamp as BridgingClampDP from deepmd.dpmodel.descriptor.dpa4c import DescrptDPA4C as DescrptDPA4CDP from deepmd.pt_expt.common import ( + register_dpmodel_mapping, torch_module, ) from deepmd.pt_expt.descriptor.base_descriptor import ( @@ -81,6 +83,21 @@ def _promote_trainable_tree(module: torch.nn.Module) -> torch.nn.Module: return module +# The distance clamp of the bridging window is parameter-free (the two window +# fractions only, no serialize()) and belongs to DPA4C alone; rebuild it fresh +# from the stored constructor arguments. +@torch_module +class BridgingClamp(BridgingClampDP): + def forward(self, *args: Any, **kwargs: Any) -> Any: + return self.call(*args, **kwargs) + + +register_dpmodel_mapping( + BridgingClampDP, + lambda v: BridgingClamp(v.f_inner, v.f_outer), +) + + @BaseDescriptor.register("dpa4c") @torch_module class DescrptDPA4C(DescrptDPA4CDP): @@ -102,6 +119,11 @@ class DescrptDPA4C(DescrptDPA4CDP): _update_sel_cls = UpdateSel + #: The per-element length scales of the bridging window follow from + #: ``type_map``, which the constructor resolves on every load, so the + #: buffer they land in stays out of the state dict. + CONFIG_DERIVED_ARRAYS = ("contact_radius",) + def adam_route_patterns(self) -> list[str]: """ Name patterns, relative to the descriptor, of the tensors that take the @@ -433,22 +455,69 @@ def _set_compression( and the cutoff survive without rounding, and the angular coupling layout keeps ``int32``; every other artifact is the ``float32`` the kernel consumes. + + The operator scalars are that metadata followed by the two fractions of + the bridging window, from which the kernels resolve the switch on the + edge envelope and the clamped length the table reads after scaling + them by the length scale of each pair. The fractions are descriptor + configuration rather than table metadata, and equal fractions denote a + descriptor without a window. + + The per-element length scales those fractions are measured against join + the artifacts in the same ``float32`` layout, so that an operator call + hands the kernels a frozen table instead of recasting one. They follow + ``type_map``, which the constructor resolves on every load, so the + buffer they land in stays out of the state dict. """ + from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( + contact_radius_input, + ) + device = self.stddev.device info = torch.as_tensor(artifacts["info"]) - self._compression_scalars = tuple( - float(value) for value in info.detach().cpu().tolist() + window = ( + (0.0, 0.0) + if self.bridging_switch is None + else (self.bridging_f_inner, self.bridging_f_outer) + ) + self._compression_scalars = ( + *(float(value) for value in info.detach().cpu().tolist()), + *window, ) for name in self._COMPRESSION_BUFFER_NAMES: dtype = self._COMPRESSION_BUFFER_DTYPES.get(name, torch.float32) value = artifacts[name].to(device=device, dtype=dtype).contiguous() - buffer_name = f"compress_{name}" - if buffer_name in self._buffers: - self._buffers[buffer_name] = value - else: - self.register_buffer(buffer_name, value) + self._store_compression_buffer(name, value) + self._store_compression_buffer( + "contact_radius", contact_radius_input(self), persistent=False + ) self.compress = True + def _store_compression_buffer( + self, + name: str, + value: torch.Tensor, + persistent: bool = True, + ) -> None: + """Bind one ``compress_``-prefixed artifact, replacing any earlier one. + + Parameters + ---------- + name + Artifact name without the ``compress_`` prefix. + value + Tensor to bind, already in the layout the kernels read. + persistent + Whether the buffer belongs in the state dict. Artifacts that follow + from the configuration alone are rebuilt on every load and are + bound without it. + """ + buffer_name = f"compress_{name}" + if buffer_name in self._buffers: + self._buffers[buffer_name] = value + else: + self.register_buffer(buffer_name, value, persistent=persistent) + def apply_charge_state(self, charge_spin: Any) -> None: """Re-specialize a compressed snapshot to a frame charge state. @@ -595,6 +664,7 @@ def fused_energy_force_graph( atom_bias: torch.Tensor, do_atomic_virial: bool, spin: torch.Tensor | None = None, + pair_table: torch.Tensor | None = None, ) -> ( tuple[ torch.Tensor, @@ -611,6 +681,8 @@ def fused_energy_force_graph( Returns ``None`` when the model or graph cannot use the level-two CUDA path, allowing the caller to retain the generic autograd lower. The trailing output is the magnetic force, empty for a spin-free model. + ``pair_table`` carries the analytical pair potential of a bridged + model, which the backward edge scan evaluates alongside the descriptor. """ if ( self.training @@ -648,4 +720,5 @@ def fused_energy_force_graph( atype.shape[0], do_atomic_virial, None if self.spin is None else self.require_spin(spin), + pair_table, ) diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 7a381b4938..e4ff2f813d 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -14,6 +14,9 @@ Any, ) +from deepmd.dpmodel.atomic_model.inner_potential import ( + InnerPotentialAtomicModel, +) from deepmd.dpmodel.train import ( AbstractTrainEntrypoint, TrainEntrypointOptions, @@ -34,6 +37,9 @@ from deepmd.pt_expt.utils.lmdb_dataset import ( LmdbDataSystem, ) +from deepmd.utils.data import ( + min_pair_dist_requirement, +) from deepmd.utils.data_system import ( DeepmdDataSystem, get_data, @@ -817,13 +823,34 @@ def change_bias( mock_loss = get_loss({"inference": True}, 1.0, len(type_map), model_to_change) data.add_data_requirements(mock_loss.label_requirement) data.add_data_requirements(get_additional_data_requirement(model_to_change)) + # A bridged model draws its statistics under the frame filter of its + # window, as training does. The filter is read from the training + # configuration; an archive that carries none cannot be filtered and + # is refused rather than read unfiltered. + if model_params is None: + children = getattr(model_to_change.atomic_model, "models", ()) + if any(isinstance(child, InnerPotentialAtomicModel) for child in children): + raise ValueError( + f"'{input_file}' carries no training configuration, so the " + "frame filter of its bridging window cannot be derived; run " + "change-bias on the checkpoint (.pt) instead." + ) + pair_filter = None + else: + pair_filter = min_pair_dist_requirement(model_params, 0.0) + if pair_filter is not None: + data.add_data_requirements([pair_filter]) if numb_batch != 0: nbatches = numb_batch else: # Cap at the minimum across systems so no system wraps and # overweights short systems (matching PT behavior). nbatches = min(data.get_nbatches()) - sampled_data = make_stat_input(data, nbatches) + sampled_data = make_stat_input( + data, + nbatches, + 0.0 if pair_filter is None else float(pair_filter.default), + ) model_to_change = model_change_out_bias( model_to_change, sampled_data, _bias_adjust_mode=bias_adjust_mode ) diff --git a/deepmd/pt_expt/fitting/ener_fitting.py b/deepmd/pt_expt/fitting/ener_fitting.py index bd3faf28ac..6d18384143 100644 --- a/deepmd/pt_expt/fitting/ener_fitting.py +++ b/deepmd/pt_expt/fitting/ener_fitting.py @@ -43,6 +43,7 @@ def call_graph( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, vacuum_descriptor: torch.Tensor | None = None, + node_gate: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Graph-native fitting forward, fused when the backend supports it. @@ -53,7 +54,9 @@ def call_graph( device rather than a traced tensor, because every export traces on CPU and moves the program afterwards. A vacuum reference is a per-type constant on the eligible configuration and enters the operator through - the bias. + the bias; the readout gate of a bridged descriptor scales the + deviation of the fused output from the bias of the type (see + :meth:`readout_reference`). """ if ( not self.training @@ -68,7 +71,14 @@ def call_graph( atom_bias = atom_bias - self.vacuum_property(vacuum_descriptor).to( atom_bias.dtype ) - return graph_fitting(self, descriptor, atype, atom_bias) + ret = graph_fitting(self, descriptor, atype, atom_bias) + if node_gate is not None: + energy = ret[self.var_name] + reference = self.readout_reference().to(energy.dtype)[atype] + ret[self.var_name] = reference + node_gate.to(energy.dtype)[:, None] * ( + energy - reference + ) + return ret return EnergyFittingNetDP.call_graph( self, descriptor, @@ -79,4 +89,5 @@ def call_graph( fparam=fparam, aparam=aparam, vacuum_descriptor=vacuum_descriptor, + node_gate=node_gate, ) diff --git a/deepmd/pt_expt/kernels/cuda/dpa4/so2_conv.py b/deepmd/pt_expt/kernels/cuda/dpa4/so2_conv.py index de407189ed..f53ec55e26 100644 --- a/deepmd/pt_expt/kernels/cuda/dpa4/so2_conv.py +++ b/deepmd/pt_expt/kernels/cuda/dpa4/so2_conv.py @@ -17,9 +17,9 @@ focus-stream count, an attention layout matching the value stream, two or more mixing layers with an identity final layer, and a radial mixer that is either absent or ``degree_channel`` of any rank. The kernels are templated on degree -and focus width only; every other dimension is a runtime argument. The -bridging-mode source gate reshapes the softmax normalization and is declined at -call time. +and focus width only; every other dimension is a runtime argument. A bridged +model reaches the kernels unchanged, because its source gate is folded into the +edge envelope the kernels already weight by. Usage and pitfalls ------------------ diff --git a/deepmd/pt_expt/kernels/dpa4c/canonical.py b/deepmd/pt_expt/kernels/dpa4c/canonical.py index 59ef8a25bc..750bb31b25 100644 --- a/deepmd/pt_expt/kernels/dpa4c/canonical.py +++ b/deepmd/pt_expt/kernels/dpa4c/canonical.py @@ -37,10 +37,13 @@ def canonical_model_eligible(model: Any) -> bool: if backend_device_type() != "cuda": return False atomic_model = getattr(model, "atomic_model", None) - descriptor = getattr(atomic_model, "descriptor", None) - fitting = getattr(atomic_model, "fitting_net", None) - if descriptor is None or fitting is None: + # A bridged composition deploys like its learned part: the fused operator + # evaluates the analytical pair potential in the same edge scan. + parts = None if atomic_model is None else atomic_model.fused_decomposition() + if parts is None: return False + descriptor = parts[0].descriptor + fitting = parts[0].fitting_net from deepmd.pt_expt.descriptor.dpa4c import ( DescrptDPA4C, ) @@ -96,12 +99,15 @@ def _forward_fake( spin: torch.Tensor, spin_pair: torch.Tensor, spin_type: torch.Tensor, + contact_radius: torch.Tensor, lmax: int, table_stride: float, table_max: float, rcut: float, eps: float, degree_floor: float, + f_inner: float, + f_outer: float, ) -> tuple[torch.Tensor, torch.Tensor]: del ( source, @@ -117,11 +123,14 @@ def _forward_fake( output_inv_std, spin_pair, spin_type, + contact_radius, table_stride, table_max, rcut, eps, degree_floor, + f_inner, + f_outer, ) from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( descriptor_profile, @@ -206,24 +215,23 @@ def _cpu_energy_gradient( *args: Any, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Reference sequence of the fused operator, evaluated in one run.""" - from deepmd.pt_expt.kernels.graph_fitting import _cpu_backward as fitting_backward - from deepmd.pt_expt.kernels.graph_fitting import _cpu_forward as fitting_forward - descriptor_args = args[:_DESCRIPTOR_ARGUMENT_COUNT] - ws, bs, resnets, w_head, b_head, bias_atom_e, act, seed, _tile = args[ + ws, bs, resnets, w_head, b_head, bias_atom_e, act, seed, _tile, pair_table = args[ _DESCRIPTOR_ARGUMENT_COUNT: ] descriptor, state = _cpu_forward(*descriptor_args) atype = descriptor_args[3] - energy, saved = fitting_forward( + energy, saved = torch.ops.deepmd.graph_fitting( descriptor, atype, ws, bs, resnets, w_head, b_head, bias_atom_e, act ) - gradient = fitting_backward( + gradient = torch.ops.deepmd.graph_fitting_backward( seed.reshape(-1, 1), saved, ws, bs, resnets, w_head, act ) - edge_gradient, spin_gradient, edge_spin_gradient = _cpu_backward( - gradient, state, *descriptor_args + edge_gradient, spin_gradient, edge_spin_gradient, pair_energy = _generic_backward( + gradient, state, *descriptor_args, pair=(pair_table, seed) ) + if pair_table.numel(): + energy = energy + pair_energy return energy, edge_gradient, spin_gradient, edge_spin_gradient @@ -257,12 +265,12 @@ def _generic_topology( # The compact ABI drops the three topology tensors of the generic ABI and its # leading ``canonical`` flag; the remaining trailing scalars are identical. -_CANONICAL_SCALAR_COUNT = 6 +_CANONICAL_SCALAR_COUNT = 8 #: Leading arguments of the fused operator that describe the descriptor: #: ``edge_vec`` plus the compact topology, the compression artifacts and the -#: six trailing geometry scalars. -_DESCRIPTOR_ARGUMENT_COUNT = 23 +#: eight trailing geometry scalars. +_DESCRIPTOR_ARGUMENT_COUNT = 26 def _cpu_forward(*args: Any) -> tuple[torch.Tensor, torch.Tensor]: @@ -288,7 +296,15 @@ def _cpu_forward(*args: Any) -> tuple[torch.Tensor, torch.Tensor]: ) -def _cpu_backward(*args: Any) -> torch.Tensor: +def _generic_backward( + *args: Any, + pair: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Run the generic reference backward on a compact payload. + + ``pair`` is the series table and the energy cotangent of an analytical + pair potential; the compact descriptor operators carry none. + """ from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( _reference_backward as generic_backward, ) @@ -300,6 +316,7 @@ def _cpu_backward(*args: Any) -> torch.Tensor: source, destination_row_ptr, ) + no_pair = edge_vec.new_empty(0) return generic_backward( descriptor_gradient, state, @@ -312,13 +329,20 @@ def _cpu_backward(*args: Any) -> torch.Tensor: *tail[:-_CANONICAL_SCALAR_COUNT], True, *tail[-_CANONICAL_SCALAR_COUNT:], + *((no_pair, no_pair) if pair is None else pair), ) -def _cpu_backward_inplace(*args: Any) -> torch.Tensor: - edge_gradient = _cpu_backward(*args) +def _cpu_backward(*args: Any) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return _generic_backward(*args)[:3] + + +def _cpu_backward_inplace( + *args: Any, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + gradients = _cpu_backward(*args) args[1].zero_() - return edge_gradient + return gradients def ensure_registered() -> None: @@ -364,6 +388,7 @@ def dpa4c_canonical_compress_energy_force( atom_bias: torch.Tensor, do_atomic_virial: bool, spin: torch.Tensor | None = None, + pair_table: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -399,6 +424,10 @@ def dpa4c_canonical_compress_energy_force( spin Per-node magnetic moments with shape ``(N, 3)`` for a spin-conditioned descriptor, or ``None``. + pair_table + Series constants of an analytical pair potential with shape + ``((T + 1) ** 2, 8)``, or ``None``. The fused operator adds its energy + to the node energies and its slope to the edge gradient. Returns ------- @@ -474,6 +503,7 @@ def dpa4c_canonical_compress_energy_force( network.activation, ownership.to(torch.float64).reshape(-1).contiguous(), node_tile(), + graph.edge_vec.new_empty(0) if pair_table is None else pair_table, ) atom_energy = atom_energy_raw * ownership[:, None].to(atom_energy_raw.dtype) energy = frame_scalar_sum(atom_energy, graph.n_node) diff --git a/deepmd/pt_expt/kernels/dpa4c/graph_compress.py b/deepmd/pt_expt/kernels/dpa4c/graph_compress.py index 3599528d20..a7805d5c1a 100644 --- a/deepmd/pt_expt/kernels/dpa4c/graph_compress.py +++ b/deepmd/pt_expt/kernels/dpa4c/graph_compress.py @@ -25,6 +25,18 @@ Degrees three and four carry a single channel each, so their couplings are driven by a compact sparse Cartesian Gaunt table rather than by specialized code. + +A bridged model closes the envelope from the inside with the inner bridging +switch of the true pair length and reads the radial table at the clamped +length, frozen at the window midpoint; the edge scan resolves both from the +pair's window, and the backward folds the table cotangent through the slope +of the clamped length. A bridged model also adds an analytical pair potential +to the learned energy. Its energy and radial slope share their exponentials +and need no saved state, so the backward operator evaluates both in the edge +scan that emits the edge gradient: given the per-pair series table and the +energy cotangent of every node, it adds the pair slope to the radial cotangent +of each edge and returns the pair energy of each node. The descriptor forward +is unaware of the term. """ from __future__ import ( @@ -59,9 +71,15 @@ backend_device_type, operator_available, ) +from deepmd.utils.bridging import ( + window_midpoint, +) from deepmd.utils.charge_state import ( validate_charge_state, ) +from deepmd.utils.element_radii import ( + UNIT_CONTACT_RADIUS, +) if TYPE_CHECKING: from collections.abc import ( @@ -1060,14 +1078,51 @@ def _table_lookup( ) -def _c3_envelope(radius: torch.Tensor, rcut: float) -> torch.Tensor: - """Evaluate the fixed exponent-five DPA4 C³ envelope.""" +def _cutoff_factor(radius: torch.Tensor, rcut: float) -> torch.Tensor: + """Evaluate the exponent-five C³ cutoff factor.""" u = ((float(rcut) - radius) / float(rcut)).clamp(0.0, 1.0) x = 1.0 - u series = 1.0 + x * (4.0 + x * (10.0 + x * (20.0 + 35.0 * x))) return u**4 * series +def _edge_envelope( + radius: torch.Tensor, + contact: torch.Tensor, + rcut: float, + f_inner: float, + f_outer: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Evaluate the edge envelope and the length the radial table reads. + + The window of a pair is two dimensionless fractions of that pair's own + length scale ``contact``, so both are functions of the reduced distance + ``radius / contact`` through the window parameter ``t``. The envelope is + the exponent-five C³ cutoff factor of the clamped length times the inner + bridging switch ``h(t)`` of the true one; the clamped length is frozen at + the window midpoint below the inner fraction and rises with the + antiderivative ``S(t)`` of the switch, ``S(1) = 1/2``, to meet the true + length at the outer fraction. Equal fractions denote a model without a + bridging window, whose table reads the true length and whose envelope is + the cutoff factor alone. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor] + The envelope and the length the table reads, both with the shape of + ``radius``. + """ + if f_outer <= f_inner: + return _cutoff_factor(radius, rcut), radius + width = f_outer - f_inner + t = ((radius / contact - f_inner) / width).clamp(0.0, 1.0) + switch = t**4 * (35.0 + t * (-84.0 + t * (70.0 - 20.0 * t))) + rise = t**5 * (7.0 + t * (-14.0 + t * (10.0 - 2.5 * t))) + clamped = contact * (window_midpoint(f_inner, f_outer) + width * rise) + seen = torch.where(radius >= f_outer * contact, radius, clamped) + return _cutoff_factor(seen, rcut) * switch, seen + + def _half_gram(value: torch.Tensor) -> torch.Tensor: """Return the Frobenius-isometric upper-triangular channel Gram.""" width = value.shape[-1] @@ -1136,6 +1191,7 @@ def _reference_descriptor( coupling_value: torch.Tensor, output_mean: torch.Tensor, output_inv_std: torch.Tensor, + contact_radius: torch.Tensor, canonical: bool, lmax: int, table_stride: float, @@ -1143,6 +1199,8 @@ def _reference_descriptor( rcut: float, eps: float, degree_floor: float, + f_inner: float, + f_outer: float, *, spin: torch.Tensor | None = None, spin_pair: torch.Tensor | None = None, @@ -1150,6 +1208,10 @@ def _reference_descriptor( ) -> torch.Tensor: """Reference implementation of the compressed DPA4C descriptor. + ``contact_radius`` holds the per-element length scale of every type-table + row in Å; the sum of the two entries of a pair is the distance the window + fractions ``f_inner`` and ``f_outer`` are measured against. + The native spin block is optional. It comprises the raw per-node moment ``spin`` with shape ``(N, 3)``, the ordered scale and shift cache ``spin_pair`` with shape ``((T + 1) ** 2, C_s, 2)``, and the per-type @@ -1173,12 +1235,17 @@ def _reference_descriptor( neighbor_type = atype[source] mask = edge_mask & (neighbor_type < type_count - 1) & (center_type < type_count - 1) maskf = mask.to(compute.dtype) - envelope = _c3_envelope(radius, float(rcut)) * maskf + radii = contact_radius.to(compute.dtype) + contact = radii[center_type] + radii[neighbor_type] + envelope, seen = _edge_envelope( + radius, contact, float(rcut), float(f_inner), float(f_outer) + ) + envelope = envelope * maskf # === Step 2. Evaluate the ordered FiLM amplitude === tabulated = _table_lookup( table.to(compute.device), - radius, + seen, float(table_stride), float(table_max), channels + radial_modes, @@ -1547,6 +1614,51 @@ def _closed_form_222_coordinate(profile: DescriptorProfile) -> int: # === Custom-operator registration === +def _reference_pair_energy( + edge_vec: torch.Tensor, + edge_index: torch.Tensor, + edge_mask: torch.Tensor, + atype: torch.Tensor, + pair_table: torch.Tensor, + eps: float, +) -> torch.Tensor: + r"""Evaluate the analytical pair energy of every node. + + Each valid edge :math:`j\to i` adds half of the screened Coulomb series + :math:`V_{ab}(\rho)=\rho^{-1}\sum_k A_{abk}e^{-c_{abk}\rho}` of its ordered + type pair to node :math:`i`, at the regularized distance the kernels use. + Rows of the padding type vanish, so padding nodes contribute nothing. + + Parameters + ---------- + edge_vec + Edge vectors with shape ``(E, 3)`` in Å. + edge_index + ``[source, destination]`` node indices with shape ``(2, E)``. + edge_mask + Valid-edge mask with shape ``(E,)``. + atype + Node types with shape ``(N,)``. + pair_table + Series constants ``[A_1..A_4, c_1..c_4]`` with shape + ``((T + 1) ** 2, 8)``. + eps + Distance regularization in Å. + + Returns + ------- + torch.Tensor + Pair energies with shape ``(N, 1)`` in eV, fp64. + """ + source, destination = edge_index[0].long(), edge_index[1].long() + type_count = math.isqrt(int(pair_table.shape[0])) + radius = torch.sqrt(edge_vec.double().square().sum(dim=-1) + eps * eps) + row = pair_table.double()[atype[destination] * type_count + atype[source]] + series = (row[:, :4] * torch.exp(-row[:, 4:] * radius[:, None])).sum(dim=-1) + energy = torch.zeros(atype.shape[0], dtype=torch.float64, device=edge_vec.device) + return energy.index_add(0, destination, 0.5 * series / radius * edge_mask)[:, None] + + def _reference_forward(*args: Any) -> tuple[torch.Tensor, torch.Tensor]: """CPU custom-op implementation returning descriptor and opaque state.""" descriptor = _reference_descriptor( @@ -1558,7 +1670,7 @@ def _reference_forward(*args: Any) -> tuple[torch.Tensor, torch.Tensor]: ) profile = descriptor_profile( int(args[9].shape[1]), - int(args[20]), + int(args[21]), args[16].ndim == 2, ) state = torch.zeros( @@ -1590,6 +1702,7 @@ def _forward_fake( spin: torch.Tensor, spin_pair: torch.Tensor, spin_type: torch.Tensor, + contact_radius: torch.Tensor, canonical: bool, lmax: int, table_stride: float, @@ -1597,6 +1710,8 @@ def _forward_fake( rcut: float, eps: float, degree_floor: float, + f_inner: float, + f_outer: float, ) -> tuple[torch.Tensor, torch.Tensor]: del ( edge_index, @@ -1614,12 +1729,15 @@ def _forward_fake( output_inv_std, spin_pair, spin_type, + contact_radius, canonical, table_stride, table_max, rcut, eps, degree_floor, + f_inner, + f_outer, ) profile = descriptor_profile( int(type_embedding.shape[1]), int(lmax), spin.ndim == 2 @@ -1639,17 +1757,23 @@ def _forward_fake( return descriptor, state +#: Position of the pair table among the backward inputs that follow +#: ``edge_vec``: nineteen tensors, ``canonical``, ``lmax`` and seven geometry +#: scalars precede it, and the pair seed closes the list. +_PAIR_TABLE_SLOT = 28 + + def _backward_fake( descriptor_gradient: torch.Tensor, state: torch.Tensor, edge_vec: torch.Tensor, *args: Any, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: del descriptor_gradient, state - spin = args[15] + atype, spin, pair_table = args[4], args[15], args[_PAIR_TABLE_SLOT] has_spin = spin.ndim == 2 - # Each absent output is allocated separately: the schema declares three - # unannotated results, so two of them may not share storage. + # Each absent output is allocated separately: the schema declares four + # unannotated results, so no two of them may share storage. return ( torch.empty_like(edge_vec), spin.new_empty((spin.shape[0], 3)) @@ -1658,6 +1782,10 @@ def _backward_fake( torch.empty_like(edge_vec) if has_spin else edge_vec.new_empty((0,), dtype=torch.float32), + edge_vec.new_empty( + (atype.shape[0], 1) if pair_table.numel() else (0,), + dtype=torch.float64, + ), ) @@ -1666,7 +1794,7 @@ def _reference_backward( state: torch.Tensor, edge_vec: torch.Tensor, *args: Any, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """CPU custom-op backward returning the coordinate and magnetic cotangents. The device operator splits the magnetic cotangent because a @@ -1675,23 +1803,41 @@ def _reference_backward( differentiates the whole node axis at once, so it returns the complete cotangent on the node axis and leaves the per-edge part at zero, which the same reduction carries through unchanged. + + A non-empty pair table adds the analytical pair potential: its energy, + weighted by the per-node cotangent ``pair_seed``, joins the differentiated + objective, so the coordinate cotangent carries the slope of the pair term, + and the node energies of the term are returned as the fourth result. """ del state spin = args[15] + pair_table, pair_seed = args[_PAIR_TABLE_SLOT:] has_spin = spin.ndim == 2 value = edge_vec.detach().clone().requires_grad_(True) moment = spin.detach().clone().requires_grad_(has_spin) + pair_energy = edge_vec.new_empty((0,), dtype=torch.float64) with torch.enable_grad(): descriptor = _reference_descriptor( value, *args[:15], - *args[18:], + *args[18:_PAIR_TABLE_SLOT], spin=moment, spin_pair=args[16], spin_type=args[17], ) + objective = (descriptor * descriptor_gradient.to(descriptor.dtype)).sum() + if pair_table.numel(): + pair_energy = _reference_pair_energy( + value, + edge_index=args[0], + edge_mask=args[1], + atype=args[4], + pair_table=pair_table, + eps=float(args[24]), + ) + objective = objective + (pair_energy[:, 0] * pair_seed.reshape(-1)).sum() gradients = torch.autograd.grad( - (descriptor * descriptor_gradient.to(descriptor.dtype)).sum(), + objective, (value, moment) if has_spin else (value,), ) return ( @@ -1700,6 +1846,7 @@ def _reference_backward( torch.zeros_like(edge_vec) if has_spin else edge_vec.new_empty((0,), dtype=torch.float32), + pair_energy.detach(), ) @@ -1711,8 +1858,8 @@ def _reference_backward( def _setup_context(ctx: Any, inputs: tuple, output: tuple) -> None: - ctx.save_for_backward(output[1], *inputs[:19]) - ctx.scalars = inputs[19:] + ctx.save_for_backward(output[1], *inputs[:20]) + ctx.scalars = inputs[20:] ctx.mark_non_differentiable(output[1]) ctx.set_materialize_grads(False) @@ -1744,13 +1891,18 @@ def _backward( "`deepmd.pt_expt.kernels.dpa4c.graph_compress.dpa4c_graph_compress`, " "which supplies it." ) + # The descriptor alone is differentiated here, so the backward receives + # no analytical pair potential. + no_pair = descriptor_gradient.new_empty(0) edge_gradient = torch.ops.deepmd.dpa4c_graph_compress_backward( descriptor_gradient, tensors[0], *tensors[1:], *ctx.scalars, + no_pair, + no_pair, )[0] - return (edge_gradient,) + (None,) * 25 + return (edge_gradient,) + (None,) * 28 @cache @@ -1776,6 +1928,40 @@ def ensure_registered() -> None: _register_ops() +def contact_radius_input(descriptor: Any) -> torch.Tensor: + """Build the per-element length scales in the form the operator consumes. + + The kernels read them only inside a bridging window, but the operator ABI + carries one entry per type-table row unconditionally. An unbridged + descriptor, which keeps no table of its own, therefore supplies the unit + table, whose pair sums are exactly 1 Å. + + A compressed descriptor freezes the result alongside its other artifacts, + so an operator call on the deployment path reads the buffer rather than + building the table again. + + Parameters + ---------- + descriptor + pt_expt DPA4C descriptor. + + Returns + ------- + torch.Tensor + Length scales in Å with shape ``(ntypes + 1,)``, contiguous fp32 on + the descriptor's device. + """ + radii = descriptor.contact_radius + if radii is None: + return torch.full( + (descriptor.ntypes + 1,), + UNIT_CONTACT_RADIUS, + dtype=torch.float32, + device=descriptor.stddev.device, + ) + return radii.to(dtype=torch.float32).contiguous() + + def compressed_operator_arguments( descriptor: Any, spin: torch.Tensor | None = None, @@ -1795,7 +1981,8 @@ def compressed_operator_arguments( ------- tuple Radial table, ordered caches, readout projections, coupling tables, - output calibration, and the native spin block. + output calibration, the native spin block, and the per-element length + scales of the bridging window. """ empty = descriptor.compress_spin_type.new_empty(0) return ( @@ -1812,6 +1999,7 @@ def compressed_operator_arguments( empty if spin is None else spin.to(torch.float32).contiguous(), descriptor.compress_spin_pair, descriptor.compress_spin_type, + descriptor.compress_contact_radius, ) @@ -1890,15 +2078,19 @@ def forward( @staticmethod def backward(ctx: Any, descriptor_gradient: torch.Tensor) -> tuple: state, edge_vec, source_order, source_row_ptr = ctx.saved_tensors + no_pair = descriptor_gradient.new_empty(0) ( edge_gradient, spin_gradient, edge_spin_gradient, + _pair_energy, ) = torch.ops.deepmd.dpa4c_graph_compress_backward( descriptor_gradient.contiguous(), state, edge_vec, *ctx.operator_args, + no_pair, + no_pair, ) spin_gradient = spin_gradient + reduce_edge_spin_gradient( edge_spin_gradient, @@ -1983,6 +2175,7 @@ def dpa4c_graph_compress_energy_force( node_capacity: int, do_atomic_virial: bool, spin: torch.Tensor | None = None, + pair_table: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -2013,6 +2206,10 @@ def dpa4c_graph_compress_energy_force( Whether to return per-node virials. spin Per-node magnetic moments with shape ``(N, 3)``, or ``None``. + pair_table + Series constants of an analytical pair potential with shape + ``((T + 1) ** 2, 8)``, or ``None``. The backward scan adds its energy + to the node energies and its slope to the edge gradient. Returns ------- @@ -2040,6 +2237,9 @@ def dpa4c_graph_compress_energy_force( from deepmd.pt_expt.kernels.edge_force_virial import ( ensure_registered as ensure_force_registered, ) + from deepmd.pt_expt.kernels.edge_force_virial import ( + frame_scalar_sum, + ) from deepmd.pt_expt.kernels.graph_fitting import ( ensure_registered as ensure_fitting_registered, ) @@ -2075,25 +2275,36 @@ def dpa4c_graph_compress_energy_force( *operator_args, ) - energy, atom_energy, descriptor_gradient = fitting_energy_and_gradient( + atom_energy, descriptor_gradient = fitting_energy_and_gradient( fitting, node_descriptor, atype, ownership, atom_bias, - graph.n_node, ) del node_descriptor + # The energy cotangent of the pair term is the ownership of each node; + # empty tensors denote a model without an analytical pair potential. + no_pair = edge_vec.new_empty(0) + seed = no_pair if pair_table is None else ownership.to(torch.float64) ( edge_gradient, spin_gradient, edge_spin_gradient, + pair_energy, ) = torch.ops.deepmd.dpa4c_graph_compress_backward( descriptor_gradient, state, edge_vec, *operator_args, + no_pair if pair_table is None else pair_table, + seed, ) + if pair_table is not None: + # The fitting energies are masked by ownership already, whereas the + # scan returns the pair energy of every node. + atom_energy = atom_energy + pair_energy * seed[:, None] + energy = frame_scalar_sum(atom_energy, graph.n_node) # The on-site magnetic gradient closes in the node kernel; the neighbour # part belongs to source nodes, and the force assembly already walks that # grouping, so it is reduced there rather than in a pass of its own. @@ -2121,8 +2332,7 @@ def fitting_energy_and_gradient( atype: torch.Tensor, ownership: torch.Tensor, atom_bias: torch.Tensor, - n_node_per_frame: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor]: """Evaluate the fused fitting network and its descriptor cotangent. Parameters @@ -2137,21 +2347,14 @@ def fitting_energy_and_gradient( Boolean mask selecting energy-contributing nodes with shape ``(N,)``. atom_bias Combined atomic energy bias with shape ``(ntypes,)``. - n_node_per_frame - Node count of each frame with shape ``(F,)``. Returns ------- - energy - Per-frame energy with shape ``(F, 1)``, fp64. atom_energy Per-node energy with shape ``(N, 1)``, fp64. descriptor_gradient Cotangent of the invariant descriptor with shape ``(N, D)``. """ - from deepmd.pt_expt.kernels.edge_force_virial import ( - frame_scalar_sum, - ) from deepmd.pt_expt.kernels.graph_fitting import ( energy_and_input_gradient, ) @@ -2164,5 +2367,4 @@ def fitting_energy_and_gradient( atom_bias, ) atom_energy = atom_energy_raw * ownership[:, None].to(atom_energy_raw.dtype) - energy = frame_scalar_sum(atom_energy, n_node_per_frame) - return energy, atom_energy, descriptor_gradient + return atom_energy, descriptor_gradient diff --git a/deepmd/pt_expt/model/dp_linear_model.py b/deepmd/pt_expt/model/dp_linear_model.py index 85cc927f4e..d158353511 100644 --- a/deepmd/pt_expt/model/dp_linear_model.py +++ b/deepmd/pt_expt/model/dp_linear_model.py @@ -1,5 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import types +from collections.abc import ( + Callable, +) from typing import ( Any, ) @@ -47,6 +50,14 @@ class LinearEnergyModel(DPModelCommon, DPLinearModel_): forward_lower_graph_exportable_with_comm = ( EnergyModel.forward_lower_graph_exportable_with_comm ) + # The compact canonical deployment resolves its descriptor, fitting net + # and analytical pair potential through the atomic model's fused + # decomposition, so a bridged compressed composition deploys through the + # same two methods as its learned part alone. + forward_lower_canonical_graph = EnergyModel.forward_lower_canonical_graph + forward_lower_canonical_graph_exportable = ( + EnergyModel.forward_lower_canonical_graph_exportable + ) def __init__( self, @@ -56,6 +67,12 @@ def __init__( DPModelCommon.__init__(self) DPLinearModel_.__init__(self, *args, **kwargs) + def compute_fitting_input_stats( + self, sampled_func: Callable[[], list[dict]] + ) -> None: + """Recompute the input statistics of every learned sub-model's fitting net.""" + self.atomic_model.compute_fitting_input_stats(sampled_func) + def forward( self, coord: torch.Tensor, diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index c02236d9c4..f7e021a684 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -24,6 +24,7 @@ from .make_model import ( _translate_energy_keys, + fused_atom_bias, make_model, ) from .model import ( @@ -206,9 +207,10 @@ def forward_lower_canonical_graph( atype, edge_vec, ) - descriptor = self.atomic_model.descriptor - fitting = self.atomic_model.fitting_net - atom_bias = fitting.bias_atom_e[:, 0] + self.atomic_model.out_bias[0, :, 0] + learned, pair_potential = self.atomic_model.fused_decomposition() + descriptor = learned.descriptor + fitting = learned.fitting_net + atom_bias = fused_atom_bias(self.atomic_model, learned) if use_dpa4c: energy, atom_energy, force, virial, atom_virial, force_mag = ( dpa4c_canonical_compress_energy_force( @@ -220,6 +222,7 @@ def forward_lower_canonical_graph( atom_bias, do_atomic_virial, spin, + None if pair_potential is None else pair_potential.pair_table, ) ) else: diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index c87ff016b5..2071875b5c 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -293,9 +293,9 @@ def _dpa4_family_child_builder(sub: dict) -> "BaseModel | None": the semantics of a standalone ``type: "dpa4"`` model -- the descriptor/fitting type defaults, the exclusion consistency check, and the loud rejections of unsupported options (``lora``, ``use_compile``) - -- instead of the generic component build that - would silently ignore them. Returns ``None`` for non-DPA4-family - children so the shared builder uses its generic path. + -- instead of the generic component build that would silently ignore + them. Returns ``None`` for non-DPA4-family children so the shared + builder uses its generic path. Parameters ---------- @@ -319,11 +319,10 @@ def get_linear_model(model_params: dict) -> BaseModel: Children with a ``descriptor`` build as learned atomic models; ``pairtab`` children build as pair-tabulation atomic models; an - ``inner_potential`` child builds the analytical bridging term, with - the learned sibling descriptor's ``inner_clamp_r_inner``/``_outer`` - derived from the child's ``r_inner``/``r_outer`` (issue #5948). A - top-level ``spin`` section (scheme ``native``) wraps the composition - as a :class:`NativeSpinEnergyModel`. + ``inner_potential`` child builds the analytical bridging term, with the + learned sibling descriptor's window derived from the child's (issue + #5948). A top-level ``spin`` section (scheme ``native``) wraps the + composition as a :class:`NativeSpinEnergyModel`. Parameters ---------- diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index 54b3f38ca4..0dfd77ccc2 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -84,6 +84,32 @@ def _translate_energy_keys( return out +def fused_atom_bias(atomic_model: Any, learned: Any) -> torch.Tensor: + """Collect the per-type energy bias a fused pipeline adds to every atom. + + The fitting net carries its own bias and every atomic model an output + bias; a bridged composition keeps the statistics on the composition and + leaves the one of its learned part at zero, so the sum covers both forms. + + Parameters + ---------- + atomic_model : BaseAtomicModel + Atomic model of the evaluated model. + learned : DPAtomicModel + Its learned descriptor-fitting part, which is ``atomic_model`` itself + unless the model is a bridged composition. + + Returns + ------- + torch.Tensor + Energy bias with shape ``(ntypes,)`` in eV. + """ + atom_bias = learned.fitting_net.bias_atom_e[:, 0] + atomic_model.out_bias[0, :, 0] + if learned is not atomic_model: + atom_bias = atom_bias + learned.out_bias[0, :, 0] + return atom_bias + + def _fused_energy_force_graph( model: Any, graph: Any, @@ -102,23 +128,33 @@ def _fused_energy_force_graph( model dict. Returns the same keys as :func:`fit_output_to_model_output_graph`, or ``None`` when no descriptor fused path applies -- the caller then uses the level-1 autograd lower. + + The atomic model answers what the pipeline evaluates + (``fused_decomposition``): its learned descriptor-fitting part and, for a + bridged composition, the analytical pair potential that the descriptor's + edge scan evaluates alongside. Masks and graph preparation stay with the + atomic model itself, which owns the exclusions. """ am = model.atomic_model - desc = getattr(am, "descriptor", None) - fit = getattr(am, "fitting_net", None) + parts = am.fused_decomposition() + if parts is None: + return None + learned, pair_potential = parts + desc, fit = learned.descriptor, learned.fitting_net fused = getattr(desc, "fused_energy_force_graph", None) - if fused is None or fit is None: + if fused is None: return None if fit.vacuum_ref and not fit.uniform_conditioning(): # The reference varies between atoms while the operators take a # per-type bias only, so such a model uses the autograd lower. return None graph, atype, output_mask = am._prepare_graph_inputs(graph, atype) - atom_bias = fit.bias_atom_e[:, 0] + am.out_bias[0, :, 0] + atom_bias = fused_atom_bias(am, learned) if fit.vacuum_ref: # The fused operators see the real atoms alone; the vacuum reference # of every type is a constant here and enters through the bias. - atom_bias = atom_bias - fit.vacuum_property(am.vacuum_descriptor())[:, 0] + atom_bias = atom_bias - fit.vacuum_property(learned.vacuum_descriptor())[:, 0] + pair = {} if pair_potential is None else {"pair_table": pair_potential.pair_table} out = fused( fit, graph, @@ -127,6 +163,7 @@ def _fused_energy_force_graph( atom_bias, do_atomic_virial, spin, + **pair, ) if out is None: return None @@ -951,9 +988,9 @@ def _call_common_graph( "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). + # CSR pre-sort serves the compressed-DPA1 fused kernels only, + # which run for plain models alone; probe via the atomic model's + # own descriptor, which a composition does not have. _desc = getattr(self.atomic_model, "descriptor", None) with_csr = ( not self.training diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 9130f5753b..f1e7f27c58 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -65,6 +65,10 @@ normalize_batch, split_batch, ) +from deepmd.dpmodel.utils.dist_check import ( + pair_margin_frame_mask, + select_frames, +) from deepmd.dpmodel.utils.learning_rate import ( make_learning_rate_schedule, ) @@ -144,6 +148,7 @@ ) from deepmd.utils.data import ( DataRequirementItem, + min_pair_dist_requirement, ) from deepmd.utils.data_system import ( DeepmdDataSystem, @@ -1357,7 +1362,6 @@ def forward( if self._task_buf_order: try: - _fitting = self.original_model.get_fitting_net() _am = getattr(self.original_model, "atomic_model", None) _vals: list[torch.Tensor] = [] for _name in self._task_buf_order: @@ -1365,6 +1369,9 @@ def forward( _actual = _name[len(_AM_PREFIX) :] _vals.append(_am._buffers[_actual]) else: + # Only a fitting-net buffer needs the fitting net; a + # composition promotes atomic-model buffers alone. + _fitting = self.original_model.get_fitting_net() _vals.append(getattr(_fitting, _name)) task_buf_vals: tuple = tuple(_vals) except AttributeError as exc: @@ -1663,7 +1670,6 @@ def _forward_graph( if self._task_buf_order: try: - _fitting = _model.get_fitting_net() _am = getattr(_model, "atomic_model", None) _vals: list[torch.Tensor] = [] for _name in self._task_buf_order: @@ -1671,6 +1677,9 @@ def _forward_graph( _actual = _name[len(_AM_PREFIX) :] _vals.append(_am._buffers[_actual]) else: + # Only a fitting-net buffer needs the fitting net; a + # composition promotes atomic-model buffers alone. + _fitting = _model.get_fitting_net() _vals.append(getattr(_fitting, _name)) task_buf_vals: tuple = tuple(_vals) except AttributeError as exc: @@ -1900,11 +1909,28 @@ def __init__( reset_epoch_streams() self.valid_numb_batch_by_task: dict[str, int] = {} + self.min_pair_dist_by_task: dict[str, float] = {} for model_key in self.model_keys: data_requirement = list(self.losses[model_key].label_requirement) data_requirement += get_additional_data_requirement(self.models[model_key]) + task_params = ( + training_params["data_dict"][model_key] + if self.multi_task + else training_params + ) + pair_filter = min_pair_dist_requirement( + self.model_params_by_task[model_key], + float( + (task_params.get("training_data") or {}).get("min_pair_dist", 0.0) + ), + ) + self.min_pair_dist_by_task[model_key] = ( + 0.0 if pair_filter is None else float(pair_filter.default) + ) + # Only training frames are filtered, so only the training data + # pays for the derived margin. self.training_data_by_task[model_key].add_data_requirements( - data_requirement + data_requirement + ([] if pair_filter is None else [pair_filter]) ) if self.validation_data_by_task[model_key] is not None: self.validation_data_by_task[model_key].add_data_requirements( @@ -1979,9 +2005,11 @@ def __init__( @functools.lru_cache def _make_sample( - _d: DeepmdDataSystem = _data, _n: int = _nbatch + _d: DeepmdDataSystem = _data, + _n: int = _nbatch, + _min_pair_dist: float = self.min_pair_dist_by_task[model_key], ) -> list[dict[str, np.ndarray]]: - return make_stat_input(_d, _n) + return make_stat_input(_d, _n, _min_pair_dist) self._sample_funcs[model_key] = _make_sample @@ -2215,6 +2243,19 @@ def initialize_statistics( # Selective weight copy (per-branch key remapping) pretrained_state = pretrained_wrapper.state_dict() target_state = self._unwrapped.state_dict() + # A plain model and its zone-bridging composition hold the same + # learned descriptor and fitting net at different paths. + learned_paths = { + model_key: ( + _learned_part_path(self._unwrapped.model[model_key]), + _learned_part_path( + pretrained_wrapper.model[ + finetune_links[model_key].get_model_branch() + ] + ), + ) + for model_key in self.model_keys + } new_state = {} for key in target_state: if key == "_extra_state": @@ -2229,6 +2270,15 @@ def initialize_statistics( finetune_rule = finetune_links[model_key] _key_from = finetune_rule.get_model_branch() pretrained_key = key.replace(f".{model_key}.", f".{_key_from}.") + path, source_path = learned_paths[model_key] + if ( + pretrained_key not in pretrained_state + and path is not None + and source_path is not None + ): + pretrained_key = _move_learned_key( + pretrained_key, path, source_path + ) use_random = ( finetune_rule.get_random_fitting() and ".descriptor." not in key @@ -2679,7 +2729,11 @@ def get_data( if data_sys is None: return {}, {} - batch = normalize_batch(data_sys.get_batch()) + batch = ( + self._next_training_batch(data_sys, task_key) + if is_train + else normalize_batch(data_sys.get_batch()) + ) input_dict, label_dict = split_batch(batch) # Drop optional inputs whose find_* flag is False so the model sees None. @@ -2714,6 +2768,53 @@ def get_data( return input_dict, label_dict + def _next_training_batch(self, data_sys: Any, task_key: str) -> dict[str, Any]: + """Fetch the next training batch of a task, without its too-close frames. + + Frames holding an atom pair closer than the task's frame filter allows + are dropped, and a batch left empty is replaced by the next one. + The retry is local to the rank and involves no collective, so + distributed ranks stay in lockstep: every rank still contributes + exactly one batch per optimizer step. + + Parameters + ---------- + data_sys : Any + Training data system of the task. + task_key : str + Resolved task key. + + Returns + ------- + dict[str, Any] + Normalized NumPy batch with at least one frame. + + Raises + ------ + RuntimeError + If one local reader pass of consecutive batches holds no valid frame. + """ + filter_enabled = self.min_pair_dist_by_task[task_key] > 0.0 + discarded = 0 + while True: + batch = normalize_batch(data_sys.get_batch()) + frame_mask = pair_margin_frame_mask(batch, filter_enabled) + if frame_mask is None or frame_mask.all(): + return batch + if frame_mask.any(): + return select_frames(batch, frame_mask) + discarded += 1 + if discarded >= data_sys.get_batch_pass_length(): + raise RuntimeError( + f"No training frame of task {task_key!r} keeps every atom " + "pair beyond the frame-filter radius within " + f"{discarded} consecutive batches. A bridged model filters " + "at the midpoint of its bridging window, so lower " + "`bridging_fraction_inner` or `bridging_fraction_outer` " + "there, or `training_data.min_pair_dist` on any other " + "model, or clean the dataset." + ) + def _epoch_length(self, model_key: str) -> int: """Return the steps this rank takes during one epoch of a task. @@ -2735,7 +2836,9 @@ def _epoch_length(self, model_key: str) -> int: length ``ceil(max_i(nbatches[i] / sys_probs[i]))``. LMDB data reports that global count while its sampler shards batches evenly across ranks; legacy data systems remain replicated. In both cases one rank takes - ``ceil(total / world_size)`` steps per epoch. + ``ceil(total / world_size)`` optimizer steps per scheduled epoch. This + differs from the local reader pass length for replicated data, which + the data system exposes through ``get_batch_pass_length``. """ data = self.training_data_by_task[model_key] total = compute_total_numb_batch(data.nbatches, data.sys_probs) @@ -3437,6 +3540,33 @@ def _more_loss_to_float(cls, more_loss: dict[str, Any]) -> dict[str, float]: } +def _learned_part_path(model: torch.nn.Module) -> str | None: + """Return the state-dict path of the learned descriptor-fitting part of a model. + + A plain model holds that part as its atomic model and a composition with + one learned sub-model, such as a zone-bridging one, as that sub-model, so + fine-tuning between the two transfers the part across the paths. A model + without a single learned part returns ``None``. + """ + atomic_model = model.atomic_model + if getattr(atomic_model, "descriptor", None) is not None: + return "atomic_model." + learned = [ + index + for index, child in enumerate(getattr(atomic_model, "models", ())) + if getattr(child, "descriptor", None) is not None + ] + return f"atomic_model.models.{learned[0]}." if len(learned) == 1 else None + + +def _move_learned_key(key: str, path: str, source_path: str) -> str: + """Rewrite a descriptor or fitting-net key from one learned-part path to another.""" + for part in ("descriptor.", "fitting_net."): + if f".{path}{part}" in key: + return key.replace(f".{path}{part}", f".{source_path}{part}", 1) + return key + + def model_change_out_bias( _model: Any, _sample_func: Any, diff --git a/deepmd/pt_expt/utils/inner_potential.py b/deepmd/pt_expt/utils/inner_potential.py index 7e7eb9284a..52f7ea3407 100644 --- a/deepmd/pt_expt/utils/inner_potential.py +++ b/deepmd/pt_expt/utils/inner_potential.py @@ -20,10 +20,9 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: return self.call(*args, **kwargs) -# InnerPotential 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). +# InnerPotential carries no trainable state (only the per-pair series table, +# derived from the constructor arguments), so it implements no +# serialize()/deserialize(); rebuild it fresh from (type_map, mode). register_dpmodel_mapping( InnerPotentialDP, lambda v: InnerPotential(type_map=v.type_map, mode=v.mode), diff --git a/deepmd/pt_expt/utils/lmdb_dataset.py b/deepmd/pt_expt/utils/lmdb_dataset.py index f16de39a01..9483b3ef66 100644 --- a/deepmd/pt_expt/utils/lmdb_dataset.py +++ b/deepmd/pt_expt/utils/lmdb_dataset.py @@ -232,6 +232,10 @@ def get_nsystems(self) -> int: """Return one logical LMDB training dataset.""" return 1 + def get_batch_pass_length(self) -> int: + """Return the local pass length belonging to the last returned batch.""" + return self._batch_iterator.batch_pass_length + @property def nbatches(self) -> list[int]: """Return the global batch count of one full pass.""" diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index aeb11d798a..b3edf999ad 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -2432,7 +2432,9 @@ def _charge_state_descriptor(data: dict, metadata: dict) -> Any | None: ) model = BaseModel.deserialize(data["model"]) - descriptor = getattr(getattr(model, "atomic_model", None), "descriptor", None) + # A bridged composition carries its charge state in its learned part. + parts = model.atomic_model.fused_decomposition() + descriptor = None if parts is None else parts[0].descriptor if ( descriptor is None or not getattr(descriptor, "compress", False) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 982c5d2a26..9f9c98c26c 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -32,6 +32,10 @@ from deepmd.utils.argcheck_nvnmd import ( nvnmd_args, ) +from deepmd.utils.bridging import ( + DEFAULT_FRACTION_INNER, + DEFAULT_FRACTION_OUTER, +) from deepmd.utils.eval_metrics import ( ENERGY_FULL_VALIDATION_PROFILE, FULL_VALIDATION_PROFILES, @@ -3639,7 +3643,7 @@ def model_args( "For an energy model it is the energy of an isolated atom, so with `vacuum_ref` in the fitting net an atom " "without neighbors gives exactly this energy. Four forms are accepted: a dict keyed by element symbol, " "e.g. `{'energy': {'O': -430.1, 'H': -13.6}}`; the name of a bundled table of isolated-atom energies, e.g. " - "`{'energy': 'omat24'}`, one of `omat24`, `omol25`, `omc25` and `odac25`, which takes precedence " + "`{'energy': 'omat24'}`, one of `mptraj`, `omat24`, `omol25`, `omc25` and `odac25`, which takes precedence " "over a file of the same name; " "the path of a JSON file holding such a dict, e.g. `{'energy': 'e0.json'}`, relative to the working directory; " "or a list with one entry per type of the `type_map` (`null` leaves a type unassigned), which is also the form " @@ -3849,25 +3853,44 @@ def standard_model_args() -> Argument: def _bridging_method_args() -> list[Argument]: """The concise analytical-bridging arguments, shared by the model types - that accept the ``bridging_method`` sugar (``dpa4`` and ``standard``). + that accept the ``bridging_method`` sugar (``dpa4`` and ``standard``, the + latter with a DPA4C descriptor). """ doc_bridging_method = ( - "Short-range bridging method. Currently supports 'ZBL'. " + "Short-range bridging method. Supports 'ZBL' and 'NLH'. " "The value is case-insensitive; set it to 'None' to disable bridging. " "This concise form is the recommended interface; it expands to the " "equivalent explicit `linear_ener` composition over the learned " "model and an `inner_potential` sub-model." ) + doc_bridging_fraction_inner = ( + "Inner radius of the bridging window, as a fraction of the covalent bond " + "length of the atom pair it applies to. A pair closer than this interacts " + "through the analytical potential alone. In DPA4 the distance the " + "descriptor reads for the pair is frozen just below the window midpoint, " + "the two atoms disappear from the descriptors of every other atom, and their " + "fitting output falls to the bias of their types, which under `vacuum_ref` " + "leaves them the energy of isolated atoms; in DPA4C the pair is removed " + "from the descriptor. Because the window follows each pair's own size, " + "one value describes every element combination. Only used when " + "`bridging_method` is enabled and no explicit radius is given." + ) + doc_bridging_fraction_outer = ( + "Outer radius of the bridging window, as a fraction of the covalent bond " + "length of the atom pair it applies to. Pairs beyond it reach the learned " + "descriptor unchanged, and the transition zone between the two fractions " + "uses a C^3-continuous septic Hermite polynomial. Only used when " + "`bridging_method` is enabled and no explicit radius is given." + ) doc_bridging_r_inner = ( - "Inner clamping radius in Å. ML descriptor distances below this radius are frozen. " - "Only used when `bridging_method` is enabled. " - "For ZBL bridging, set `training.training_data.min_pair_dist` to the same value " - "so frames with atom pairs closer than `bridging_r_inner` are skipped during training." + "Inner radius of the bridging window in Å, applied to every atom pair alike. " + "Giving it together with `bridging_r_outer` replaces the fractional window " + "of `bridging_fraction_inner` and `bridging_fraction_outer`." ) doc_bridging_r_outer = ( - "Outer clamping radius in Å. The transition zone " - "`[bridging_r_inner, bridging_r_outer]` uses a C^3-continuous " - "septic Hermite polynomial. Only used when `bridging_method` is enabled." + "Outer radius of the bridging window in Å, applied to every atom pair alike. " + "Giving it together with `bridging_r_inner` replaces the fractional window " + "of `bridging_fraction_inner` and `bridging_fraction_outer`." ) return [ Argument( @@ -3877,18 +3900,32 @@ def _bridging_method_args() -> list[Argument]: default="None", doc=supported_backends("pt", "pt_expt") + doc_bridging_method, ), + Argument( + "bridging_fraction_inner", + float, + optional=True, + default=DEFAULT_FRACTION_INNER, + doc=supported_backends("pt", "pt_expt") + doc_bridging_fraction_inner, + ), + Argument( + "bridging_fraction_outer", + float, + optional=True, + default=DEFAULT_FRACTION_OUTER, + doc=supported_backends("pt", "pt_expt") + doc_bridging_fraction_outer, + ), Argument( "bridging_r_inner", float, optional=True, - default=0.5, + default=None, doc=supported_backends("pt", "pt_expt") + doc_bridging_r_inner, ), Argument( "bridging_r_outer", float, optional=True, - default=0.8, + default=None, doc=supported_backends("pt", "pt_expt") + doc_bridging_r_outer, ), ] @@ -4119,34 +4156,57 @@ def inner_potential_model_args() -> Argument: injected only into the ``linear_ener`` ``models`` variant. """ doc_mode = ( - "The analytical pair-potential formula. Currently supports 'zbl' " - "(case-insensitive)." + "The analytical pair-potential formula. Supports 'zbl', the ZBL " + "universal screened repulsion, and 'nlh', a repulsion refitted to " + "self-consistent density-functional pair energies (case-insensitive)." + ) + doc_fraction_inner = ( + "Inner radius of the bridging window, as a fraction of the covalent bond " + "length of the atom pair it applies to, derived onto the learned sibling's " + "descriptor: a closer atom pair interacts through the analytical potential " + "alone. Because the window follows each pair's own size, one value " + "describes every element combination." + ) + doc_fraction_outer = ( + "Outer radius of the bridging window, as a fraction of the covalent bond " + "length of the atom pair it applies to. The transition zone between the " + "two fractions uses a C^3-continuous septic Hermite polynomial." ) doc_r_inner = ( - "Inner clamping radius in Å, applied to the learned sibling's " - "descriptor: ML descriptor distances below this radius are frozen. " - "For ZBL bridging, set `training.training_data.min_pair_dist` to the " - "same value so frames with atom pairs closer than `r_inner` are " - "skipped during training." + "Inner radius of the bridging window in Å, applied to every atom pair " + "alike. Giving it together with `r_outer` replaces the fractional window." ) doc_r_outer = ( - "Outer clamping radius in Å, applied to the learned sibling's " - "descriptor. The transition zone `[r_inner, r_outer]` uses a " - "C^3-continuous septic Hermite polynomial." + "Outer radius of the bridging window in Å, applied to every atom pair " + "alike. Giving it together with `r_inner` replaces the fractional window." ) ca = Argument( "inner_potential", dict, [ Argument("mode", str, optional=True, default="zbl", doc=doc_mode), - Argument("r_inner", float, optional=True, default=0.5, doc=doc_r_inner), - Argument("r_outer", float, optional=True, default=0.8, doc=doc_r_outer), + Argument( + "fraction_inner", + float, + optional=True, + default=DEFAULT_FRACTION_INNER, + doc=doc_fraction_inner, + ), + Argument( + "fraction_outer", + float, + optional=True, + default=DEFAULT_FRACTION_OUTER, + doc=doc_fraction_outer, + ), + Argument("r_inner", float, optional=True, default=None, doc=doc_r_inner), + Argument("r_outer", float, optional=True, default=None, doc=doc_r_outer), ], doc=supported_backends("pt", "pt_expt") + "Analytical short-range bridging pair potential (e.g. ZBL), usable " - "only as a sub-model of a `linear_ener` composition; the clamping " - "radii are derived onto the learned sibling's descriptor at build " - "time.", + "only as a sub-model of a `linear_ener` composition over one DPA4 or " + "DPA4C model; the window radii are derived onto the learned sibling's " + "descriptor at build time.", ) return ca @@ -5746,11 +5806,16 @@ def training_data_args() -> list[ "Frames containing any atom pair closer than this distance are excluded " "from loss computation, as DFT labels for near-collision configurations " "are often unreliable. Set to 0 to disable (default). " - "Under distributed training (DDP/FSDP), if any rank has no valid frame " - "in its current batch, every rank collectively skips that training step. " - "Note: enabling this adds an O(N²) distance check per frame in the " - "DataLoader workers (CPU-side), which may slow down training for large " - "systems. To avoid the overhead, consider pre-cleaning the dataset instead." + "A bridged model filters its training frames by the midpoint of its own " + "bridging window, below which its learned energy is throttled to less " + "than half of its amplitude; that radius is taken from the model, so this " + "key must be left unset for such a model and is an error otherwise. " + "A batch left without a valid frame is never trained on: under " + "distributed training the pt backend skips that step on every rank, and " + "the pt_expt backend draws the next batch on the rank concerned. " + "Note: enabling this adds a geometric check per frame on the CPU side " + "of the data pipeline, which may slow down training for large systems. " + "To avoid the overhead, consider pre-cleaning the dataset instead." ) args = [ @@ -5795,7 +5860,7 @@ def training_data_args() -> list[ float, optional=True, default=0.0, - doc=supported_backends("pt") + doc_min_pair_dist, + doc=supported_backends("pt", "pt_expt") + doc_min_pair_dist, ), ] diff --git a/deepmd/utils/bridging.py b/deepmd/utils/bridging.py index fbdc3c103d..b5bdb2ea94 100644 --- a/deepmd/utils/bridging.py +++ b/deepmd/utils/bridging.py @@ -10,7 +10,7 @@ "models": [ {"type": "dpa4", "descriptor": {...}, "fitting_net": {...}}, {"type": "inner_potential", "mode": "zbl", - "r_inner": 0.5, "r_outer": 0.8} + "fraction_inner": 0.26, "fraction_outer": 0.80} ] } @@ -24,14 +24,56 @@ import copy +import numpy as np + __all__ = [ + "check_bridging_record_version", + "check_window_inside_cutoff", + "concise_bridging_window", "expand_bridging_method", "is_bridged_sezm_config", + "migrate_inner_clamp_keys", + "resolve_bridging_window", + "window_midpoint", ] +BRIDGING_RECORD_VERSION = 1.3 +"""Descriptor format version that introduced the current window mechanics.""" + +DEFAULT_FRACTION_INNER = 0.26 +"""Inner window radius as a fraction of the pair's covalent bond length.""" + +DEFAULT_FRACTION_OUTER = 0.80 +"""Outer window radius as a fraction of the pair's covalent bond length.""" + _DPA4_FAMILY_TYPES = ("dpa4", "sezm") +def window_midpoint(f_inner: float, f_outer: float) -> float: + """ + Midpoint of a bridging window, in the units of its two radii. + + The switch of the window is exactly half open there, the training-frame + filter of a bridged model keeps pairs down to it, and the DPA4C clamp holds + that same value below the inner radius. Reading the point from one place + keeps the descriptors and the data filter agreeing on where the training + data end. + + Parameters + ---------- + f_inner : float + Inner radius of the window. + f_outer : float + Outer radius of the window. + + Returns + ------- + float + The midpoint ``(f_inner + f_outer) / 2``. + """ + return 0.5 * (f_inner + f_outer) + + def is_bridged_sezm_config(data: dict) -> bool: """Return whether a config is a bridged DPA4/SeZM linear composition. @@ -84,6 +126,8 @@ def _is_dpa4_family(sub: dict) -> bool: # composition nor the learned child. _CONSUMED_KEYS = ( "bridging_method", + "bridging_fraction_inner", + "bridging_fraction_outer", "bridging_r_inner", "bridging_r_outer", ) @@ -238,6 +282,32 @@ def route_canonical_learned_options(composition: dict, learned: dict) -> None: learned[key] = copy.deepcopy(composition[key]) +def concise_bridging_window(data: dict) -> dict: + """Read the window the concise ``bridging_*`` keys of a model section spell. + + Parameters + ---------- + data : dict + A model section carrying the concise bridging keys. + + Returns + ------- + dict + The window fields of an ``inner_potential`` child: the two fractions, + defaulted, and the two explicit radii, ``None`` when unset. + """ + return { + "fraction_inner": float( + data.get("bridging_fraction_inner", DEFAULT_FRACTION_INNER) + ), + "fraction_outer": float( + data.get("bridging_fraction_outer", DEFAULT_FRACTION_OUTER) + ), + "r_inner": data.get("bridging_r_inner"), + "r_outer": data.get("bridging_r_outer"), + } + + def expand_bridging_method(data: dict) -> dict: """Expand the ``bridging_method`` sugar into a ``linear_ener`` config. @@ -284,8 +354,7 @@ def expand_bridging_method(data: dict) -> dict: "and an `inner_potential` sub-model instead." ) data = copy.deepcopy(data) - r_inner = float(data.get("bridging_r_inner", 0.5)) - r_outer = float(data.get("bridging_r_outer", 0.8)) + window = concise_bridging_window(data) # Legacy promotion (pt `type: "dpa4"` semantics): a descriptor-scoped # exclusion also governs the analytical term of a bridged model. @@ -314,8 +383,7 @@ def expand_bridging_method(data: dict) -> dict: { "type": "inner_potential", "mode": method, - "r_inner": r_inner, - "r_outer": r_outer, + **window, }, ], "atom_exclude_types": data.get("atom_exclude_types", []), @@ -327,3 +395,153 @@ def expand_bridging_method(data: dict) -> dict: if key in data: canonical[key] = data[key] return canonical + + +def resolve_bridging_window(inner_cfg: dict) -> dict: + """ + Resolve an ``inner_potential`` child's window into descriptor options. + + The window is a pair of fractions of each atom pair's own covalent bond + length, so one setting describes every element combination. Explicit radii + in Å replace it with a window that is the same for every pair; both forms + reach the descriptor as the same three options, the absolute one through a + unit length scale that leaves the fractions carrying the radii themselves. + + Parameters + ---------- + inner_cfg : dict + The ``inner_potential`` sub-model configuration. + + Returns + ------- + dict + The ``inner_clamp_f_inner``, ``inner_clamp_f_outer`` and + ``inner_clamp_scale`` options of the learned sibling's descriptor. + + Raises + ------ + ValueError + If only one of ``r_inner`` and ``r_outer`` is given. + """ + r_inner = inner_cfg.get("r_inner") + r_outer = inner_cfg.get("r_outer") + if (r_inner is None) != (r_outer is None): + raise ValueError( + "An explicit bridging window needs both `r_inner` and `r_outer`; " + "leave both unset to size the window by the covalent bond length " + "of each atom pair." + ) + if r_inner is None: + return { + "inner_clamp_f_inner": float( + inner_cfg.get("fraction_inner", DEFAULT_FRACTION_INNER) + ), + "inner_clamp_f_outer": float( + inner_cfg.get("fraction_outer", DEFAULT_FRACTION_OUTER) + ), + "inner_clamp_scale": "covalent", + } + return { + "inner_clamp_f_inner": float(r_inner), + "inner_clamp_f_outer": float(r_outer), + "inner_clamp_scale": "absolute", + } + + +def migrate_inner_clamp_keys(config: dict) -> None: + """Bring a descriptor configuration's window options up to date in place. + + Records predating the pair-relative window spell it as ``inner_clamp_r_inner`` + and ``inner_clamp_r_outer``, two radii in Å. The absolute scale is that same + window measured against a unit pair length, so the radii carry over as the + fractions themselves and the record needs no other adjustment. + + Parameters + ---------- + config : dict + The constructor arguments read back from a serialized descriptor. + """ + r_inner = config.pop("inner_clamp_r_inner", None) + r_outer = config.pop("inner_clamp_r_outer", None) + if r_inner is None and r_outer is None: + return + # A record that states only one radius is carried through as it stands, so + # the constructor reports the half-given window rather than this function. + config["inner_clamp_f_inner"] = None if r_inner is None else float(r_inner) + config["inner_clamp_f_outer"] = None if r_outer is None else float(r_outer) + config["inner_clamp_scale"] = "absolute" + + +def check_window_inside_cutoff( + contact_radius: np.ndarray, f_outer: float, rcut: float +) -> None: + """Refuse a bridging window that reaches the cutoff for some pair. + + The switch and the clamp of a pair return to the identity at the outer + radius. A pair whose outer radius lies at or beyond the cutoff leaves the + neighbor list before that happens, so its gate never reopens and the + displayed distance of its edges is read past the cutoff. + + Parameters + ---------- + contact_radius : np.ndarray + Per-type length scales in Å with shape ``(ntypes + 1,)``, the trailing + entry belonging to the padding type. + f_outer : float + Outer radius of the window as a fraction of the pair length scale. + rcut : float + Cutoff radius in Å. + + Raises + ------ + ValueError + If the outer radius of the largest pair reaches ``rcut``. + """ + outer = f_outer * 2.0 * float(np.max(contact_radius[:-1])) + if outer >= rcut: + raise ValueError( + f"The bridging window of the largest pair ends at {outer:.3f} Å, at or " + f"beyond the cutoff {rcut} Å; enlarge `rcut` or shrink the window." + ) + + +def check_bridging_record_version(config: dict, version: float) -> None: + """Refuse legacy state when the receiving descriptor enables bridging. + + Bridged records before :data:`BRIDGING_RECORD_VERSION` froze the clamp at the + inner radius, muted a closing pair through a full product of switch amplitudes + on the messages and carried no readout gate, so their weights were trained + under a different function of the geometry than the one this code builds. + State dictionaries do not encode whether the source used bridging, so the + guard uses the receiving descriptor's configuration. It also rejects a + legacy plain state loaded directly into a bridged descriptor. An unbridged + receiving descriptor passes this guard; its version migrations still apply. + + Parameters + ---------- + config : dict + The receiving descriptor's constructor arguments, with the window keys + in their current spelling. These come from the record for portable + deserialization, or from the target model for state-dict loading. + version : float + The stored descriptor version, from ``@version`` or ``version_tensor``. + + Raises + ------ + ValueError + If the receiving descriptor is bridged and the state predates the + current window mechanics. + """ + if ( + version < BRIDGING_RECORD_VERSION + and config.get("inner_clamp_f_inner") is not None + ): + raise ValueError( + f"Cannot load DPA4/SeZM state at format version {version} into a bridged " + f"descriptor; bridging requires version {BRIDGING_RECORD_VERSION} or " + "later. State dictionaries do not identify whether the source used " + "bridging. For a plain version-1.1 or version-1.2 checkpoint, use " + "--finetune to migrate the source in its original unbridged model " + "before enabling bridging. If the source was already bridged, its " + "window mechanics are incompatible. Retrain the bridged model." + ) diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index a232aab393..4d73e58d87 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -29,6 +29,11 @@ LRU_CACHE_SIZE, ) from deepmd.utils import random as dp_random +from deepmd.utils.bridging import ( + expand_bridging_method, + resolve_bridging_window, + window_midpoint, +) from deepmd.utils.path import ( DPH5Path, DPPath, @@ -125,6 +130,10 @@ def __init__( self.type_idx_map, np.array([-1], dtype=np.int32) ) self.type_map = type_map + # Element-based derived fields may use the model's names without a + # dataset map. Those names do not define the reader's numeric type + # domain: legacy spin data also stores unnamed virtual atom types. + self._element_type_map = tuple(type_map or self.type_map or ()) if type_map is None and self.type_map is None and self.mixed_type: raise RuntimeError("mixed_type format must have type_map!") # make idx map @@ -167,6 +176,7 @@ def add( dtype: np.dtype | None = None, output_natoms_for_type_sel: bool = False, special_shape: str | None = None, + length_scale: str = "absolute", ) -> "DeepmdData": """Add a data item that to be loaded. @@ -198,6 +208,12 @@ def add( special_shape : str, optional Name of a loader-defined non-standard shape contract. ``"hessian"`` stores one full-frame ``(3 * natoms) x (3 * natoms)`` matrix per frame. + length_scale : str, optional + Per-element length scale a derived geometric item measures its + default against: ``"absolute"`` gives every element the unit radius, + so the default is a distance in Å, while ``"covalent"`` gives each + element its covalent radius, so the default is a fraction of each + atom pair's own bond length. """ # normalize key: "atomic_" prefix -> "atom_", same convention as _load_set output if key.startswith("atomic_"): @@ -213,6 +229,7 @@ def add( "default": default, "dtype": dtype, "output_natoms_for_type_sel": output_natoms_for_type_sel, + "length_scale": length_scale, } if special_shape is not None: # Preserve the established dictionary schema for ordinary labels; @@ -327,8 +344,124 @@ def get_batch(self, batch_size: int) -> dict: idx = np.arange(self.iterator, iterator_1, dtype=np.int64) self.iterator += batch_size ret = self._get_subdata(self.batch_set, idx) + if "pair_margin" in self.data_dict: + ret["find_pair_margin"] = np.float32(1.0) + ret["pair_margin"] = self._derive_pair_margin_batch( + ret["coord"], + ret["box"] if self.pbc else None, + ret["type"], + ) return ret + def _derive_pair_margin( + self, + coord: np.ndarray, + box: np.ndarray | None, + atype: np.ndarray, + ) -> np.ndarray: + """Derive the `pair_margin` requirement of one frame. + + The requirement has no file behind it: the frame and the batch read + paths compute it from the frame geometry. The value is the frame's + margin against the window each atom pair carries, so the filter keeps + a frame whose value reaches one, and the scan stops at the first pair + that settles the question. + + Parameters + ---------- + coord + Coordinates of one frame with shape (natoms * 3,) in Å. + box + Cell of one frame with shape (9,) in Å, or None without periodicity. + atype + Atom types of one frame with shape (natoms,). + + Returns + ------- + np.ndarray + Minimum pair margin with shape (1,), dimensionless. + """ + from deepmd.dpmodel.utils.dist_check import ( + compute_min_pair_margin_single, + ) + + return np.array( + [ + compute_min_pair_margin_single( + coord, + box, + atype, + self._pair_margin_half_thresholds(), + screened=True, + ) + ], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ) + + def _pair_margin_half_thresholds(self) -> np.ndarray: + """Resolve the per-type half-thresholds of the `pair_margin` window. + + Returns + ------- + np.ndarray + Half-thresholds in Å with shape (ntypes,), read-only because the + table is shared by every reader of the same window and type map. + """ + from deepmd.dpmodel.utils.dist_check import ( + requirement_half_thresholds, + ) + + requirement = self.data_dict["pair_margin"] + return requirement_half_thresholds( + float(requirement.get("default", 0.0)), + str(requirement.get("length_scale", "absolute")), + self._element_type_map or None, + max(self.get_ntypes(), len(self._element_type_map)), + ) + + def _derive_pair_margin_batch( + self, + coord: np.ndarray, + box: np.ndarray | None, + atype: np.ndarray, + ) -> np.ndarray: + """Derive the `pair_margin` requirement of a whole batch. + + The batch is scanned in one pass, which is what a batch read pays for + the requirement instead of one scan per frame. The value settles each + frame against the window the requirement declares, as + :meth:`_derive_pair_margin` does, and equals the minimum pair + margin whenever that margin decides the frame; a frame the window + accepts may carry a larger margin that it reaches. + + Parameters + ---------- + coord + Coordinates of the batch with shape (nframes, natoms * 3) in Å. + box + Cells of the batch with shape (nframes, 9) in Å, or None without + periodicity. + atype + Atom types of the batch with shape (nframes, natoms). + + Returns + ------- + np.ndarray + Margin of every frame, with shape (nframes, 1), dimensionless. + """ + from deepmd.dpmodel.utils.dist_check import ( + compute_min_pair_margin_batch, + ) + + margins = compute_min_pair_margin_batch( + coord, + box, + atype, + self._pair_margin_half_thresholds(), + screened=True, + ) + return margins.reshape(-1, 1).astype(GLOBAL_NP_FLOAT_PRECISION, copy=False) + def get_test(self, ntests: int = -1) -> dict: """Get the test data with `ntests` frames. @@ -556,24 +689,13 @@ def get_single_frame(self, index: int, num_worker: int) -> dict: frame_data["fid"] = index - # === Compute min_pair_dist on-the-fly in DataLoader worker === - if "min_pair_dist" in self.data_dict: - from deepmd.dpmodel.utils.dist_check import ( - compute_min_pair_dist_single, - ) - - frame_data["find_min_pair_dist"] = np.float32(1.0) - min_pair_dist = float(self.data_dict["min_pair_dist"].get("default", 0.0)) - frame_data["min_pair_dist"] = np.array( - [ - compute_min_pair_dist_single( - frame_data["coord"], - frame_data.get("box"), - frame_data["type"], - stop_below=min_pair_dist, - ) - ], - dtype=GLOBAL_NP_FLOAT_PRECISION, + # === Derive the pair margin of the frame in the DataLoader worker === + if "pair_margin" in self.data_dict: + frame_data["find_pair_margin"] = np.float32(1.0) + frame_data["pair_margin"] = self._derive_pair_margin( + frame_data["coord"], + frame_data.get("box"), + frame_data["type"], ) if self.modifier is not None: @@ -1208,6 +1330,12 @@ class DataRequirementItem: available. ``"derived"`` computes the value from structural frame data. Only optional tracked fields require availability-homogeneous batching. + length_scale : {"absolute", "covalent"}, optional + Per-element length scale a derived geometric field measures its + default against. ``"absolute"`` gives every element the unit radius, + so the default is a distance in Å; ``"covalent"`` gives each element + its covalent radius, so the default is a fraction of each atom pair's + own bond length. """ def __init__( @@ -1224,6 +1352,7 @@ def __init__( output_natoms_for_type_sel: bool = False, special_shape: str | None = None, source_policy: DataRequirementSourcePolicy = "tracked", + length_scale: str = "absolute", ) -> None: if source_policy not in {"tracked", "default", "derived"}: raise ValueError( @@ -1247,6 +1376,7 @@ def __init__( self.output_natoms_for_type_sel = output_natoms_for_type_sel self.special_shape = special_shape self.source_policy = source_policy + self.length_scale = length_scale self.dict = self.to_dict() def to_dict(self) -> dict: @@ -1262,6 +1392,7 @@ def to_dict(self) -> dict: "dtype": self.dtype, "output_natoms_for_type_sel": self.output_natoms_for_type_sel, "source_policy": self.source_policy, + "length_scale": self.length_scale, } if self.special_shape is not None: data["special_shape"] = self.special_shape @@ -1288,3 +1419,107 @@ def has_data_requirement(requirements: Iterable[DataRequirementItem], key: str) of reinterpreting loss-specific configuration such as prefactors. """ return any(item.key == key for item in requirements) + + +def min_pair_dist_requirement( + model_params: dict, + min_pair_dist: float, +) -> DataRequirementItem | None: + """ + Build the frame-filter requirement a model and its dataset config imply. + + A bridged model owns the window. In its lower half the bridging modules + throttle the learned energy to less than half of its amplitude, so a label + there asks the model for a multiple of the residual it has to express, and + a frame holding a pair that close steers the fit inside the window. The + filter is therefore the window midpoint, read from the model rather than + restated in the dataset section. Any other model has no such radius, and + the dataset section may still ask for a plain distance floor. + + Parameters + ---------- + model_params : dict + The model section of the training config, in either the concise + ``bridging_method`` spelling or the canonical composition form. + min_pair_dist : float + The ``training_data.min_pair_dist`` value, a distance in Å. Zero or + less leaves an unbridged model unfiltered. + + Returns + ------- + DataRequirementItem or None + The requirement of the derived ``pair_margin`` field, or None when no + filter applies. + + Raises + ------ + ValueError + If a bridged model is given an explicit ``min_pair_dist``, which would + put a second, silently disagreeing copy of the filter radius in the + input. + """ + # A model is bridged exactly when its canonical composition carries an + # analytical child, which is the same test the composition builder applies + # and covers every learned family it accepts, DPA4C included. + canonical = expand_bridging_method(model_params) + inner = next( + ( + sub + for sub in (canonical.get("models") or []) + if isinstance(sub, dict) and sub.get("type") == "inner_potential" + ), + None, + ) + if inner is not None: + if min_pair_dist > 0.0: + raise ValueError( + "`training_data.min_pair_dist` must not be set on a bridged " + "model: the filter is the midpoint of the bridging window and " + "is taken from the model. Remove the key." + ) + window = resolve_bridging_window(inner) + fraction = window_midpoint( + window["inner_clamp_f_inner"], window["inner_clamp_f_outer"] + ) + scale = window["inner_clamp_scale"] + else: + fraction, scale = min_pair_dist, "absolute" + # An unbridged model without a floor registers no requirement, so a zero + # radius cannot leave one behind that the reader would then refuse to + # derive; a bridged window is positive by construction. + if fraction <= 0.0: + return None + return _pair_margin_requirement(fraction, scale) + + +def _pair_margin_requirement(fraction: float, length_scale: str) -> DataRequirementItem: + """Assemble the derived frame-filter requirement. + + The default carries the filter radius and the length scale names the + per-element radii it multiplies, which is everything a reader needs to + size the per-type thresholds from its own type map. The field holds each + frame's margin against those thresholds, a dimensionless ratio the filter + compares against one. + + Parameters + ---------- + fraction : float + Filter radius, as a fraction of the pair length scale. + length_scale : str + Either ``"covalent"`` or ``"absolute"``. + + Returns + ------- + DataRequirementItem + The requirement for the derived ``pair_margin`` field. + """ + return DataRequirementItem( + "pair_margin", + 1, + atomic=False, + must=False, + high_prec=False, + default=fraction, + source_policy="derived", + length_scale=length_scale, + ) diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index cea270bd37..890224e063 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -340,6 +340,7 @@ def add_dict(self, adict: dict[str, dict[str, Any]]) -> None: "output_natoms_for_type_sel", False ), special_shape=adict[kk].get("special_shape"), + length_scale=adict[kk].get("length_scale", "absolute"), ) def add_data_requirements( @@ -361,6 +362,7 @@ def add( dtype: np.dtype | None = None, output_natoms_for_type_sel: bool = False, special_shape: str | None = None, + length_scale: str = "absolute", ) -> None: """Add a data item that to be loaded. @@ -391,6 +393,9 @@ def add( If True and type_sel is True, the atomic dimension will be natoms instead of nsel special_shape : str, optional Name of a loader-defined non-standard shape contract. + length_scale : str, optional + Per-element length scale a derived geometric item measures its + default against, either ``"absolute"`` or ``"covalent"``. """ for ii in self.data_systems: ii.add( @@ -405,6 +410,7 @@ def add( dtype=dtype, output_natoms_for_type_sel=output_natoms_for_type_sel, special_shape=special_shape, + length_scale=length_scale, ) def reduce(self, key_out: str, key_in: str) -> None: @@ -650,6 +656,18 @@ def get_batch_size(self) -> int: """Get the batch size.""" return self.batch_size + def get_batch_pass_length(self) -> int: + """Return the local sampling-pass length, including system probabilities. + + NumPy data systems are replicated across ranks, so a reader's pass + length is independent of the distributed optimizer's world size. + """ + from deepmd.dpmodel.utils.training_utils import ( + compute_total_numb_batch, + ) + + return compute_total_numb_batch(self.nbatches, self.sys_probs) + def print_summary(self, name: str) -> None: print_summary( name, diff --git a/deepmd/utils/element_radii.py b/deepmd/utils/element_radii.py new file mode 100644 index 0000000000..c593f233da --- /dev/null +++ b/deepmd/utils/element_radii.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Per-element length scales used to size pair-wise geometric windows.""" + +import numpy as np + +from deepmd.utils.model_preset_data import ( + PERIODIC_TABLE, +) + +# fmt: off +COVALENT_RADII: tuple[float, ...] = ( + 0.32, 0.46, 1.33, 1.02, 0.85, 0.75, 0.71, 0.63, 0.64, 0.67, 1.55, 1.39, 1.26, + 1.16, 1.11, 1.03, 0.99, 0.96, 1.96, 1.71, 1.48, 1.36, 1.34, 1.22, 1.19, 1.16, + 1.11, 1.10, 1.12, 1.18, 1.24, 1.21, 1.21, 1.16, 1.14, 1.17, 2.10, 1.85, 1.63, + 1.54, 1.47, 1.38, 1.28, 1.25, 1.25, 1.20, 1.28, 1.36, 1.42, 1.40, 1.40, 1.36, + 1.33, 1.31, 2.32, 1.96, 1.80, 1.63, 1.76, 1.74, 1.73, 1.72, 1.68, 1.69, 1.68, + 1.67, 1.66, 1.65, 1.64, 1.70, 1.62, 1.52, 1.46, 1.37, 1.31, 1.29, 1.22, 1.23, + 1.24, 1.33, 1.44, 1.44, 1.51, 1.45, 1.47, 1.42, 2.23, 2.01, 1.86, 1.75, 1.69, + 1.70, 1.71, 1.72, 1.66, 1.66, 1.68, 1.68, 1.65, 1.67, 1.73, 1.76, 1.61, 1.57, + 1.49, 1.43, 1.41, 1.34, 1.29, 1.28, 1.21, 1.22, 1.36, 1.43, 1.62, 1.75, 1.65, + 1.57, +) +"""Molecular single-bond covalent radii in Å, indexed by atomic number minus one. + +Pyykkö and Atsumi, *Chem. Eur. J.* **15**, 186 (2009), the one published set that +covers all 118 elements. The sum of two of these radii estimates the length scale +of the corresponding pair, which is what a window expressed as a fraction of the +pair's own size is measured against. +""" +# fmt: on + +UNIT_CONTACT_RADIUS = 0.5 +"""The per-element radius that turns every pair's length scale into exactly 1 Å. + +A window given directly in Å is the fractional rule measured in this unit, so the +two forms share one implementation: the fractions then carry the radii themselves. +""" + +_SYMBOL_TO_RADIUS: dict[str, float] = dict( + zip(PERIODIC_TABLE, COVALENT_RADII, strict=True) +) + + +def covalent_radii_from_type_map(type_map: list[str]) -> np.ndarray: + """ + Look up the covalent radius of every type in a type map. + + Parameters + ---------- + type_map : list[str] + Element symbols; the index in this list is the ``atype`` value. + + Returns + ------- + np.ndarray + Radii in Å with shape ``(ntypes,)`` in float64. + + Raises + ------ + ValueError + If a symbol is not an element of the periodic table. + """ + unknown = [symbol for symbol in type_map if symbol not in _SYMBOL_TO_RADIUS] + if unknown: + raise ValueError( + f"No covalent radius is tabulated for {unknown}; every type must be " + "an element symbol." + ) + return np.asarray( + [_SYMBOL_TO_RADIUS[symbol] for symbol in type_map], dtype=np.float64 + ) + + +def contact_radius_table( + ntypes: int, type_map: list[str] | None, scale: str +) -> np.ndarray: + """ + Tabulate the length scale each type contributes to a pair. + + A pair-wise window is a fraction of ``table[a] + table[b]``. Under the + ``"covalent"`` scale that sum is the pair's covalent bond length, so one + pair of fractions describes every element combination; under ``"absolute"`` + every element carries :data:`UNIT_CONTACT_RADIUS`, the sum is 1 Å, and the + fractions are the radii themselves in Å. + + Parameters + ---------- + ntypes : int + Number of real types. + type_map : list[str], optional + Element symbols of the types, required by the ``"covalent"`` scale. + scale : str + Either ``"covalent"`` or ``"absolute"``. + + Returns + ------- + np.ndarray + Radii in Å with shape ``(ntypes + 1,)``. The trailing entry belongs to + the padding type, whose edges are masked; it carries the unit radius so + the reduced distance of a padded edge stays finite. + + Raises + ------ + ValueError + If the scale is unknown, or the covalent scale is requested without a + type map of element symbols covering every type. + """ + if scale not in ("covalent", "absolute"): + raise ValueError(f"`scale` must be 'covalent' or 'absolute', got {scale!r}.") + table = np.full(ntypes + 1, UNIT_CONTACT_RADIUS, dtype=np.float64) + if scale == "covalent": + if type_map is None or len(type_map) != ntypes: + raise ValueError( + f"A covalent window needs a `type_map` covering all {ntypes} types." + ) + table[:ntypes] = covalent_radii_from_type_map(type_map) + return table diff --git a/deepmd/utils/model_stat.py b/deepmd/utils/model_stat.py index 2ce82ace6b..109de19fd5 100644 --- a/deepmd/utils/model_stat.py +++ b/deepmd/utils/model_stat.py @@ -12,6 +12,10 @@ from deepmd.dpmodel.utils.batch import ( normalize_batch, ) +from deepmd.dpmodel.utils.dist_check import ( + pair_margin_frame_mask, + select_frames, +) log = logging.getLogger(__name__) @@ -40,6 +44,70 @@ def _get_stat_batch(data: Any, sys_idx: int) -> dict[str, Any]: return data.get_batch(sys_idx=sys_idx) +def _get_stat_pass_length(data: Any, sys_idx: int) -> int: + """Return the batch count of one pass over a statistical system.""" + get_stat_numb_batches = getattr(data, "get_stat_numb_batches", None) + if get_stat_numb_batches is not None: + return int(get_stat_numb_batches(sys_idx)) + return int(data.get_nbatches()[sys_idx]) + + +def _collect_stat_batches( + data: Any, + sys_idx: int, + nbatches: int, + min_pair_dist: float, +) -> list[dict[str, Any]]: + """Draw the normalized batches of one statistical system. + + Without a pair-clearance filter these are the next ``nbatches`` batches + of the system. With one, the frames that hold a pair inside its window are + dropped and a batch left empty is replaced by the next one, until ``nbatches`` + batches are kept or the scan ends. The scan covers at least one pass over + the system, and the system contributes whenever that pass draws a valid + frame. A pass serves every frame of an LMDB system; the NPY reader serves + whole batches and reshuffles a set before serving it again, so its pass + leaves out the frames of a set that do not fill a last batch. + + Parameters + ---------- + data + The data source, see :func:`make_stat_input`. + sys_idx : int + Index of the statistical system. + nbatches : int + Number of batches to keep. + min_pair_dist : float + Filter radius the requirement was registered with. A non-positive + value disables the filter. + + Returns + ------- + list[dict[str, Any]] + Normalized batches, each with at least one frame; empty when the + system holds no batch or no valid frame. + """ + target = _get_stat_numb_batches(data, sys_idx, nbatches) + scan_limit = target + if min_pair_dist > 0.0: + scan_limit = max(target, _get_stat_pass_length(data, sys_idx)) + batches: list[dict[str, Any]] = [] + for _ in range(scan_limit): + if len(batches) == target: + break + stat_data = _get_stat_batch(data, sys_idx) + if "natoms_vec" in stat_data: + stat_data["natoms_vec"] = stat_data["natoms_vec"].astype(np.int32) + batch = normalize_batch(stat_data) + frame_mask = pair_margin_frame_mask(batch, min_pair_dist > 0.0) + if frame_mask is not None and not frame_mask.all(): + if not frame_mask.any(): + continue + batch = select_frames(batch, frame_mask) + batches.append(batch) + return batches + + def _make_all_stat_ref(data: Any, nbatches: int) -> dict[str, list[Any]]: all_stat = defaultdict(list) for ii in range(_get_stat_nsystems(data)): @@ -57,8 +125,7 @@ def collect_batches( ) -> dict[str, list[Any]]: """Collect batches from a DeepmdDataSystem into a dict of lists. - This is a low-level helper used by the TF backend and by - :func:`make_stat_input`. + This is a low-level helper used by the TF backend. Parameters ---------- @@ -103,6 +170,7 @@ def collect_batches( def make_stat_input( data: Any, nbatches: int, + min_pair_dist: float = 0.0, ) -> list[dict[str, np.ndarray]]: """Pack data for statistics using DeepmdDataSystem. @@ -123,34 +191,47 @@ def make_stat_input( batches specifically for statistics. nbatches : int Number of batches to collect per system. + min_pair_dist : float + Radius the training-frame filter was registered with. When + positive, frames holding a pair inside its window are dropped while + the batches are drawn, a batch left empty is replaced by the next one + of its system, and a system without any valid frame is skipped, so the + statistics see the frames the optimizer sees. Returns ------- list[dict[str, np.ndarray]] Per-system dicts with concatenated numpy arrays. """ - all_stat = collect_batches(data, nbatches, merge_sys=False) - nsystems = _get_stat_nsystems(data) log.info(f"Packing data for statistics from {nsystems} systems") - keys = list(all_stat.keys()) lst: list[dict[str, np.ndarray]] = [] for ii in range(nsystems): - merged: dict[str, np.ndarray] = {} - for key in keys: - vals = all_stat[key][ii] # list of batch arrays for this system - if isinstance(vals[0], np.ndarray): - if vals[0].ndim >= 2: - merged[key] = np.concatenate(vals, axis=0) - else: - # 1D arrays (e.g. natoms_vec) — per-system constant - merged[key] = vals[0] - else: - # scalar flags like find_* - merged[key] = vals[0] - - lst.append(normalize_batch(merged)) + batches = _collect_stat_batches(data, ii, nbatches, min_pair_dist) + if not batches: + reason = ( + "no frame keeps every atom pair beyond the filter radius" + if min_pair_dist > 0.0 + else "it holds no batch" + ) + log.info(f"Skipping statistical system {ii}: {reason}.") + continue + # Arrays of two or more dimensions lead with the frame axis; every + # other entry (find_* flags, an absent box) is constant per system. + lst.append( + { + key: np.concatenate([batch[key] for batch in batches], axis=0) + if isinstance(value, np.ndarray) and value.ndim >= 2 + else value + for key, value in batches[0].items() + } + ) + if min_pair_dist > 0.0 and not lst: + raise RuntimeError( + "No sampled frame keeps every atom pair beyond the filter radius; " + "lower that radius or clean the dataset." + ) return lst diff --git a/deepmd/utils/preset_out_bias_tables.json b/deepmd/utils/preset_out_bias_tables.json index 8b2eb6b43e..b21a05811e 100644 --- a/deepmd/utils/preset_out_bias_tables.json +++ b/deepmd/utils/preset_out_bias_tables.json @@ -1,4 +1,98 @@ { + "mptraj": { + "source": "isolated-atom reference energies for MPtraj, in eV; final VASP free energy TOTEN", + "values": { + "H": -1.11735921, + "He": 0.00072954, + "Li": -0.29748611, + "Be": -0.01782239, + "B": -0.25734156, + "C": -1.25021064, + "N": -3.12441985, + "O": -1.65374417, + "F": -0.49252164, + "Ne": -0.01247796, + "Na": -0.22870161, + "Mg": -0.09610658, + "Al": -0.20483144, + "Si": -0.81475807, + "P": -1.88810233, + "S": -0.87991589, + "Cl": -0.24670388, + "Ar": -0.04953865, + "K": -0.22628721, + "Ca": -0.09267035, + "Sc": -2.21339504, + "Ti": -2.63982102, + "V": -3.73739476, + "Cr": -5.60265238, + "Mn": -5.32539558, + "Fe": -3.59646519, + "Co": -2.14803415, + "Ni": -1.05263706, + "Cu": -0.60057645, + "Zn": -0.16513938, + "Ga": -0.3492915, + "Ge": -0.89181347, + "As": -1.68361385, + "Se": -0.75776795, + "Br": -0.21051221, + "Kr": -0.03222653, + "Rb": -0.18698314, + "Sr": -0.06804051, + "Y": -2.22084712, + "Zr": -2.31307488, + "Nb": -3.1635722, + "Mo": -4.60202593, + "Tc": -3.54437113, + "Ru": -2.61646665, + "Rh": -1.72385331, + "Pd": -1.47772741, + "Ag": -0.33955104, + "Cd": -0.16845757, + "In": -0.34399097, + "Sn": -0.82535907, + "Sb": -1.4111781, + "Te": -0.64594019, + "I": -0.17802982, + "Xe": -0.0086123, + "Cs": -0.13563312, + "Ba": -0.03400297, + "La": -0.6851664, + "Ce": -1.47222392, + "Pr": -0.37210695, + "Nd": -0.25548789, + "Pm": -0.12652484, + "Sm": 0.00131425, + "Eu": -8.36636129, + "Gd": -10.36747409, + "Tb": 0.48416441, + "Dy": 0.65727533, + "Ho": -0.38379388, + "Er": -0.38610604, + "Tm": -0.31738017, + "Yb": -0.06423855, + "Lu": -0.381518, + "Hf": -3.52789166, + "Ta": -3.74334288, + "W": -4.65659151, + "Re": -4.70116882, + "Os": -2.89848439, + "Ir": -1.57731153, + "Pt": -0.56968822, + "Au": -0.28806702, + "Hg": -0.12460079, + "Tl": -0.30597824, + "Pb": -0.76492312, + "Bi": -1.32635439, + "Ac": -0.38628688, + "Th": -1.02769395, + "Pa": -2.67263716, + "U": -5.02364665, + "Np": -7.69888774, + "Pu": -10.80589773 + } + }, "odac25": { "source": "isolated-atom reference energies of the UMA odac task (ODAC25), in eV; fairchem configs/uma/training_release/element_refs/iso_atom_elem_refs.yaml (MIT license), odac_elem_refs", "values": { diff --git a/deepmd/utils/spin.py b/deepmd/utils/spin.py index 02c75b4b5f..723a7f72d3 100644 --- a/deepmd/utils/spin.py +++ b/deepmd/utils/spin.py @@ -1,5 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import copy +from collections.abc import ( + Iterator, +) +from typing import ( + Any, +) import numpy as np @@ -8,6 +14,27 @@ ) +def env_protection_descriptors( + descriptor: dict[str, Any], inherited_protection: float | None = None +) -> Iterator[dict[str, Any]]: + """Yield leaf configurations that use environment-matrix protection. + + DPA4-family descriptors regularize their geometry with ``eps`` instead of + an environment matrix. Hybrid-level protection is moved to the children + that consume it, without forwarding that option to DPA4. A child's own + value takes precedence over the value inherited from its parent. + """ + kind = descriptor["type"].lower() + if kind == "hybrid": + protection = descriptor.pop("env_protection", inherited_protection) + for child in descriptor["list"]: + yield from env_protection_descriptors(child, protection) + elif kind not in {"dpa4", "sezm", "dpa4c"}: + if inherited_protection is not None: + descriptor.setdefault("env_protection", inherited_protection) + yield descriptor + + def normalize_spin_use_spin(use_spin: list, type_map: list[str]) -> list[bool]: """Normalize ``use_spin`` to a per-type boolean list. diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index 5bb6a8097a..ad70c54b6f 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -331,7 +331,9 @@ it requires `lmax >= 1` (the default). The `deepspin` scheme augments each magnetic atom with a virtual atom at `coord + spin * virtual_scale`, which doubles the internal neighbor capacity and type map. `virtual_scale` is required by this scheme and ignored by the native -scheme. +scheme. The `deepspin` scheme cannot be combined with `bridging_method`: a +magnetic atom lies inside the bridging window of its own virtual partner, so a +bridged spin model uses the native scheme. ### Multi-task / shared fitting @@ -371,42 +373,101 @@ Best checkpoints fold the LoRA deltas back into the base weights, producing plain DPA4/SeZM checkpoints suitable for deployment. See `examples/water/dpa4/lora_ft.json`. -## Zone bridging (ZBL) +## Zone bridging -DPA4/SeZM can add an analytical short-range repulsion (typically ZBL) to the -learned energy in a protected region: +DPA4/SeZM can add an analytical short-range repulsion to the learned energy in a +protected region: ```math -E_i = E_i^{\mathrm{DPA4/SeZM}} + E_i^{\mathrm{ZBL}}. +E_i = E_i^{\mathrm{DPA4/SeZM}} + E_i^{\mathrm{pair}}. ``` -Below `bridging_r_inner` the distance seen by the descriptor is clamped, with a -smooth transition back to the true distance up to `bridging_r_outer`; a source -gate additionally blocks the learned model from leaking information about the -frozen short-range pairs. The recommended way to enable it is the concise -form, set directly on the `dpa4` model: +Two potentials are available. `bridging_method: "zbl"` uses the +Ziegler-Biersack-Littmark screened nuclear potential, which is derived from +frozen atomic electron densities and covers every element. +`bridging_method: "nlh"` uses the Nordlund-Lehtola-Hobler form, whose +coefficients are fitted to self-consistent quantum-chemical pair energies and +therefore follow the real repulsive wall much more closely, leaving a smaller +residual repulsion at ordinary bond lengths for the learned model to cancel. +`zbl` is the default; whether `nlh` yields a better trained model has not been +established. + +> [!NOTE] +> The `nlh` coefficients shipped with DeePMD-kit are this project's own refit and +> are **not** the coefficients published by Nordlund, Lehtola and Hobler. They are +> fitted independently to the same open reference data — K. Nordlund, G. Hobler and +> S. Lehtola, Zenodo [10.5281/zenodo.14172633](https://doi.org/10.5281/zenodo.14172633) +> v1.0 (2024), CC BY 4.0 — under conditions this use requires, and are therefore +> modified material. The functional form is that of +> [Phys. Rev. A **111**, 032818 (2025)](https://doi.org/10.1103/PhysRevA.111.032818). +> Pairs containing an element heavier than uranium, which the reference data does +> not cover, fall back to ZBL. + +Inside the window three things happen to a pair, all of them smooth in the +pair distance. The distance the descriptor reads for the pair is frozen two +fifths into the window for every closer pair and rejoins the true distance at +the outer radius; the freeze point sits just below the point where the training +data end, so the descriptor still resolves the pair distance where labels +exist and is asked to extrapolate by less than one percent of the window +width below them. A source gate silences every atom that has a neighbor inside +its inner radius towards every *other* atom: such an atom is absent from the +descriptors of all its other neighbors exactly as if those edges had been +deleted, so the frozen pair leaks no information into the learned energy of +the atoms around it, while the two atoms of the pair keep seeing each other at +the frozen distance. And the fitting output of an atom fades with the same +gate to the bias of its type, so a frozen atom contributes nothing but that +constant; with the +[isolated-atom energy reference](train-energy.md#isolated-atom-energy-reference) +its energy is then exactly that of an isolated atom. Along the closing of a +pair the learned energy of the pair therefore follows the switch linearly +once the descriptor's distance has frozen, and the analytical term takes over. + +The two radii follow the size of the pair they apply to. Each is a fraction of +the covalent bond length of that element pair, so one setting covers every +element combination in the type map. The defaults place the outer radius at +0.80 of the bond length, below all but the shortest multiple bonds, and the +inner radius at 0.26, deep inside the repulsive wall where the analytical term +is the responsible description. The recommended way to enable bridging is the +concise form, set directly on the `dpa4` model: ```json { "model": { "type": "dpa4", - "bridging_method": "zbl", - "bridging_r_inner": 0.5, - "bridging_r_outer": 0.8 + "bridging_method": "zbl" } } ``` -When ZBL bridging is enabled, set `training.training_data.min_pair_dist` to the -same value as `bridging_r_inner` so frames with shorter atom pairs are excluded -from training. See `examples/water/dpa4/input-zbl.json` for a complete example. +`bridging_fraction_inner` and `bridging_fraction_outer` adjust the two +fractions. Giving `bridging_r_inner` and `bridging_r_outer` instead replaces +the rule with one window in Å applied to every pair alike. See +`examples/water/dpa4/input-zbl.json` for a complete example. + +Training frames that contain a pair closer than the midpoint of its window are +left out of the loss. Below the midpoint the learned contribution of the pair is +throttled to less than half of its amplitude, so a label there would ask the +model for a multiple of the energy it has to express. The threshold is derived +from the window of each element pair, and `training_data.min_pair_dist` must +stay unset on a bridged model. The data statistics, including those of +`dp change-bias`, are drawn under the same filter, while validation batches +are not filtered, so a validation error also covers frames that training +leaves out. + +Labels inside the upper half of the window of every element pair train the +learned contribution there, and in practice they come from dimer scans that run +down to about half the covalent bond length. Below the freeze point the +learned contribution of a pair is the switch amplitude times a value the +training data fix at the frozen separation, so no label is needed to keep the +descent smooth; the atoms around a closing pair, however, see its edges fade, +and labels inside the window are what train that response. > [!NOTE] > Output-bias statistics and bridging: the model energy is -> `E = E_model + E_bias`, and the ZBL term belongs to `E_model`. The +> `E = E_model + E_bias`, and the analytical term belongs to `E_model`. The > `set-by-statistic` bias mode (initial statistics, finetune with a > random fitting, `dp change-bias --mode set`) fits `E_bias` to the raw -> data labels and by definition ignores `E_model` — the analytical ZBL +> data labels and by definition ignores `E_model` — the analytical pair > contribution included. For a self-consistent calibration of a bridged > model use `change-by-statistic`, which subtracts the complete bridged > prediction. See [change-bias](change-bias.md) for the precise @@ -438,8 +499,8 @@ form above expands to exactly this equivalent explicit form: { "type": "inner_potential", "mode": "zbl", - "r_inner": 0.5, - "r_outer": 0.8 + "fraction_inner": 0.26, + "fraction_outer": 0.8 } ] } @@ -685,8 +746,8 @@ Which models lose multi-rank, on either route: - **The deepspin (virtual-atom) spin scheme**, which overrides the export ABI to `nlist` because it expands virtual atoms inside the graph. -**ZBL zone bridging is multi-rank capable** (including combined with native -spin). The Source Freeze Propagation gate folds each node's full +**Zone bridging is multi-rank capable** (including combined with native +spin). The Source Freeze Propagation gate is a product over each node's full *outgoing*-edge set, which no single rank observes for ghost owners; the with-comm artifact completes the gate's per-node `[log eta, zero count]` partials across ranks with one reverse-accumulate plus one forward-broadcast @@ -738,7 +799,7 @@ 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 +(`add_chg_spin_ebd`) and zone bridging are graph-eligible too; only the `deepspin` virtual-atom spin scheme remains dense-only: ```bash @@ -762,8 +823,8 @@ 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 -- +of exporting a silently-dense-only artifact. `native`-scheme spin and 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 @@ -828,7 +889,7 @@ inference](#multi-gpu-mpi-inference). No native-spin combination restrictions remain on the graph route: multi-rank inference, charge-spin FiLM conditioning (`add_chg_spin_ebd`), -and ZBL zone bridging (`bridging_method: ZBL`) all combine freely with +and zone bridging (`bridging_method: zbl` or `nlh`) all combine freely with `spin.scheme: native`. ## Embedding extraction @@ -954,19 +1015,19 @@ closed over the one-hop neighbor shell. - Model compression is not supported. - Multi-rank (multi-GPU/MPI) LAMMPS inference works on both export routes: the PT `edge_vec` archive and the pt_expt NeighborGraph archive each embed - a with-comm artifact. ZBL zone bridging (and its native-spin combination) + a with-comm artifact. Zone bridging (and its native-spin combination) participates: the Source Freeze Propagation gate's per-node partials are completed across ranks by one reverse-accumulate plus one forward-broadcast border 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 + conditioning, 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 and ZBL zone bridging both combine with it, + Charge-spin FiLM conditioning and zone bridging both combine with it, multi-rank included. See [Native spin (magnetic)](#native-spin-magnetic). diff --git a/doc/model/dpa4c.md b/doc/model/dpa4c.md index 7d5fbc0ec5..7e79d8098e 100644 --- a/doc/model/dpa4c.md +++ b/doc/model/dpa4c.md @@ -202,6 +202,133 @@ and evaluation and inference are governed independently by `DP_AMP_INFER`. A model trained in full precision can therefore be evaluated under mixed precision, and the reverse. +## Zone bridging + +DPA4C can add an analytical short-range repulsion to the learned energy, so that +close approaches in high-temperature, high-pressure or collision simulations meet +a physical repulsive wall instead of an extrapolated network output: + +```math +E_i = E_i^{\mathrm{DPA4C}} + E_i^{\mathrm{pair}}. +``` + +Two potentials are available. `bridging_method: "zbl"` uses the +Ziegler-Biersack-Littmark screened nuclear potential, which is derived from +frozen atomic electron densities and covers every element. +`bridging_method: "nlh"` uses the Nordlund-Lehtola-Hobler form, whose +coefficients are fitted to self-consistent quantum-chemical pair energies and +therefore follow the real repulsive wall much more closely; it leaves a smaller +residual repulsion at ordinary bond lengths, which is energy the learned model +would otherwise have to cancel. `zbl` is the default; whether `nlh` yields a +better trained model has not been established. + +> [!NOTE] +> The `nlh` coefficients shipped with DeePMD-kit are this project's own refit and +> are **not** the coefficients published by Nordlund, Lehtola and Hobler. They are +> fitted independently to the same open reference data — K. Nordlund, G. Hobler and +> S. Lehtola, Zenodo [10.5281/zenodo.14172633](https://doi.org/10.5281/zenodo.14172633) +> v1.0 (2024), CC BY 4.0 — under conditions this use requires, and are therefore +> modified material. The functional form is that of +> [Phys. Rev. A **111**, 032818 (2025)](https://doi.org/10.1103/PhysRevA.111.032818). +> Pairs containing an element heavier than uranium, which the reference data does +> not cover, fall back to ZBL. + +The analytical term acts on every pair inside `rcut`. The learned model, in turn, is +made blind to pairs that are too close to be covered by training data. Below +the inner radius of the window a pair leaves the descriptor entirely: both +atoms are described exactly as if the other one were absent, and the pair +interacts through the analytical term alone. Between the inner and the outer +radius the pair enters the descriptor smoothly in two coupled ways: its +amplitude opens from zero to one, and the separation its radial functions +read is released at the same rate, from the midpoint of the window at the +inner radius to the true separation at the outer radius. A pair deep inside +the window therefore changes the descriptor through its amplitude alone, +with its radial features held where the retained training frames place them. +Beyond the outer radius the descriptor is the same as without bridging. The +transition keeps the energy three times continuously differentiable, like the +outer cutoff. + +The two radii follow the size of the pair they apply to. Each is a fraction of +the covalent bond length of that element pair, so a hydrogen pair is described +down to a much shorter separation than a caesium pair, and one setting covers +every element combination in the type map. The defaults place the outer radius +at 0.80 of the bond length, below all but the shortest multiple bonds, and +the inner radius at 0.26, deep inside the repulsive wall where the analytical +term is the responsible description. + +Enable it on the model: + +```json +{ + "model": { + "type_map": [ + "O", + "H" + ], + "descriptor": { + "type": "dpa4c", + "rcut": 6.0 + }, + "bridging_method": "zbl" + } +} +``` + +`type_map` must name chemical elements, because the analytical term is built +from their nuclear charges and the window from their covalent radii. `bridging_fraction_inner` and +`bridging_fraction_outer` adjust the two fractions. Giving `bridging_r_inner` +and `bridging_r_outer` instead replaces the rule with one window in Å applied +to every pair alike, which is worth doing only for a single-element system or +to reproduce an earlier model. See `examples/water/dpa4c/input-zbl.json` for a +complete input. + +Training frames that contain a pair closer than the midpoint of its window are +left out of the loss. Below the midpoint the learned contribution of the pair is +throttled to less than half of its amplitude, so a label there would ask the +model for a multiple of the energy it has to express; the separation the +radial functions read never falls below the midpoint either, and through the +lower half of the window it moves by less than 7 % of the window width, so the +radial features of a close pair stay next to the range the retained frames +cover. The threshold is derived from the window of each element pair, and +`training_data.min_pair_dist` must stay unset on a bridged model. The data +statistics, including those of `dp change-bias`, are drawn under the same +filter, while validation batches are not filtered, so a validation error also +covers frames that training leaves out. + +Bridging therefore relies on short-range training data. The labels that are +kept have to reach into the window of every element pair, which in practice +means adding dimer scans that run down to about half the covalent bond length. +Between the inner radius and the closest retained label the learned +contribution of the pair falls to zero with no data to shape it; labels inside +the window keep that descent smooth, while a pair whose window holds no label +can show spurious force maxima there. + +The same model can be written explicitly as a `linear_ener` model over the +DPA4C model and an `inner_potential` sub-model with `weights: "sum"`; the +concise form above expands to it. + +Bridging changes nothing in the workflow after training. The model freezes, +compresses and runs in LAMMPS like a plain DPA4C model: compression still +selects the compact canonical graph form, the fused kernels evaluate the +analytical term in the same pass over the neighbors as the descriptor, and multi-GPU runs +need no extra communication. Native spin and the frame charge state combine +with bridging; the analytical term depends on the separation alone and leaves +the magnetic force to the learned model. + +A model trained without bridging fine-tunes into its bridged form, and the +other way round. Add the bridging keys above to the input, keep the +architecture, and start from the earlier checkpoint: + +```bash +dp --pt-expt train input-zbl.json --finetune model.ckpt.pt +``` + +The descriptor and the fitting net are carried over although the bridged model +holds them one level deeper, inside the composition. Fine-tuning is required +rather than optional: the bridging window changes what the descriptor sees at +short range, so the learned part is trained against the analytical term it is +now summed with. + ## Model compression Compression is the deployment step. It replaces the analytic radial functions diff --git a/doc/model/train-energy.md b/doc/model/train-energy.md index f448017ae6..83bf693866 100644 --- a/doc/model/train-energy.md +++ b/doc/model/train-energy.md @@ -153,15 +153,20 @@ limit of the model is pinned. Two model options provide this together: | Name | Reference calculation | Elements | | -------- | --------------------- | ----------------------------------------- | +| `mptraj` | MPtraj | 89, H to Pu without Po, At, Rn, Fr and Ra | | `omat24` | OMat24 | 89, H to Pu without Po, At, Rn, Fr and Ra | | `omol25` | OMol25, neutral atoms | 83, H to Bi | | `omc25` | OMC25 | 94, H to Bk without Tb, Am and Cm | | `odac25` | ODAC25 | 94, H to Pu | -These four tables are the isolated-atom reference energies of the UMA training -tasks published with fairchem +The `mptraj` table uses the final VASP `free energy TOTEN`, in eV. + +The `omat24`, `omol25`, `omc25` and `odac25` tables are the isolated-atom +reference energies of the UMA training tasks published with fairchem (`configs/uma/training_release/element_refs/iso_atom_elem_refs.yaml`, MIT -license). Each table is on the energy scale of its own reference calculation +license). + +Each table is on the energy scale of its own reference calculation (the OMat24 table gives H = -1.117 eV, the OMol25 table H = -13.446 eV), so a model takes the table of the calculation that produced its training data. A table of another reference calculation is given as a JSON file diff --git a/examples/water/dpa4/input-zbl.json b/examples/water/dpa4/input-zbl.json index ddd010233c..7977151038 100644 --- a/examples/water/dpa4/input-zbl.json +++ b/examples/water/dpa4/input-zbl.json @@ -45,9 +45,7 @@ }, "use_compile": false, "enable_tf32": true, - "bridging_method": "zbl", - "bridging_r_inner": 0.5, - "bridging_r_outer": 0.8 + "bridging_method": "zbl" }, "learning_rate": { "type": "wsd", @@ -81,8 +79,7 @@ "../data/data_1", "../data/data_2" ], - "batch_size": 1, - "min_pair_dist": 0.5 + "batch_size": 1 }, "validation_data": { "systems": [ diff --git a/examples/water/dpa4c/README.md b/examples/water/dpa4c/README.md index 216e71c999..e297d322d6 100644 --- a/examples/water/dpa4c/README.md +++ b/examples/water/dpa4c/README.md @@ -14,6 +14,13 @@ The recommended new-training path selects the `Nano` grade with the named preset `descriptor` and `fitting_net`; the explicit water `type_map` replaces the element list, while `use_amp` and `seed` supply run-specific settings. +`input-zbl.json` is the same model with ZBL zone bridging: `bridging_method` +adds the analytical ZBL repulsion to the learned energy, and a pair closer than +the inner radius of the bridging window interacts through ZBL alone. The window +follows the covalent bond length of each element pair, so it needs no setting +of its own. The model trains, compresses and deploys exactly like the plain +one. + DPA4C is built for extreme-speed molecular dynamics, so its arguments are best read as a budget split between two quantities: inference throughput and the largest system that fits in memory. diff --git a/examples/water/dpa4c/input-zbl.json b/examples/water/dpa4c/input-zbl.json new file mode 100644 index 0000000000..52485370c5 --- /dev/null +++ b/examples/water/dpa4c/input-zbl.json @@ -0,0 +1,76 @@ +{ + "_comment": "DPA4C-Nano energy-training example for the water dataset with ZBL zone bridging.", + "model": { + "preset": "dpa4c-nano-v20260911", + "type_map": [ + "O", + "H" + ], + "descriptor": { + "use_amp": false, + "seed": 42 + }, + "fitting_net": { + "seed": 42 + }, + "_comment_bridging": "ZBL acts on every pair inside rcut. A pair closer than the inner radius of the window leaves the DPA4C descriptor and interacts through ZBL alone; the pair enters the descriptor smoothly up to the outer radius. Both radii are fractions of the covalent bond length of the element pair they apply to, so one setting covers every combination in the type map.", + "bridging_method": "zbl" + }, + "learning_rate": { + "type": "wsd", + "start_lr": 1e-3, + "stop_lr": 1e-6, + "warmup_ratio": 0.003, + "warmup_start_factor": 0.2, + "decay_phase_ratio": 0.2, + "decay_type": "inverse_linear" + }, + "loss": { + "type": "ener", + "loss_func": "mae", + "f_use_norm": true, + "start_pref_e": 20, + "limit_pref_e": 20, + "start_pref_f": 20, + "limit_pref_f": 20, + "start_pref_v": 5, + "limit_pref_v": 5 + }, + "optimizer": { + "type": "HybridMuon", + "weight_decay": 0.001 + }, + "training": { + "stat_file": "./dpa4c_zbl.hdf5", + "stat_file_mode": "update", + "training_data": { + "systems": [ + "../data/data_0", + "../data/data_1", + "../data/data_2" + ], + "_comment_min_pair_dist": "A bridged model derives its own frame filter from the bridging window, so `min_pair_dist` must not be set here.", + "batch_size": 1 + }, + "validation_data": { + "systems": [ + "../data/data_3" + ], + "batch_size": 1, + "numb_btch": 1 + }, + "numb_steps": 1000000, + "gradient_max_norm": 5.0, + "save_freq": 2000, + "max_ckpt_keep": 3, + "disp_file": "lcurve.out", + "disp_freq": 1000, + "disp_training": true, + "time_training": true, + "enable_compile": true, + "seed": 42 + }, + "validating": { + "save_best_dir": "ckpt_best" + } +} diff --git a/source/lmp/tests/dpa4_spin_harness.py b/source/lmp/tests/dpa4_spin_harness.py index 26a4cd3f95..3b12a4899d 100644 --- a/source/lmp/tests/dpa4_spin_harness.py +++ b/source/lmp/tests/dpa4_spin_harness.py @@ -17,7 +17,6 @@ ``test_lammps_dpa4_spin_zbl_pt2.py`` (native spin + ZBL bridging). """ -import json import os import signal import subprocess as sp @@ -37,6 +36,9 @@ from lammps import ( PyLammps, ) +from python_reference import ( + run_python_reference, +) MPI_DEFAULT_TIMEOUT = 120.0 @@ -146,7 +148,6 @@ def compute_expected( atype = (type_map - 1).tolist() # LAMMPS 1-based -> deepmd 0-based infer_dir = str(pb_file.resolve().parent) script = textwrap.dedent(f"""\ - import json import sys import numpy as np @@ -165,19 +166,15 @@ def compute_expected( atomic=True, spin=np.array({spin.tolist()!r}).reshape(1, -1, 3), ) - print(json.dumps({{ + result = {{ "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}") - # Warp and other backends may log to stdout before the final JSON payload. - result = json.loads(proc.stdout.strip().splitlines()[-1]) + result = run_python_reference(script) # Raw DeepEval force_mag (dE/dspin), scaled by LAMMPS's own # spin_norm / hbar unit convention (see HBAR_METAL) before comparison. diff --git a/source/lmp/tests/python_reference.py b/source/lmp/tests/python_reference.py new file mode 100644 index 0000000000..90028813bc --- /dev/null +++ b/source/lmp/tests/python_reference.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Isolated Python reference calculations for native LAMMPS tests.""" + +import json +import subprocess +import sys +import tempfile +from pathlib import ( + Path, +) +from typing import ( + Any, +) + + +def run_python_reference(script: str) -> Any: + """Run a reference calculation without sharing the native operator registry. + + Parameters + ---------- + script + Python code assigning its JSON-serializable return value to ``result``. + Its diagnostics may use either stdout or stderr; the result travels + through a separate file so dependency initialization cannot corrupt it. + + Returns + ------- + Any + The decoded reference value. + + Raises + ------ + RuntimeError + If the child process fails, with both diagnostic streams included. + """ + with tempfile.TemporaryDirectory(prefix="deepmd-python-reference-") as directory: + output = Path(directory) / "result.json" + program = ( + "import sys\n" + + script + + "\nimport json\n" + + "with open(sys.argv[1], 'w', encoding='utf-8') as stream:\n" + + " json.dump(result, stream)\n" + ) + proc = subprocess.run( + [sys.executable, "-c", program, str(output)], + capture_output=True, + text=True, + ) + if proc.returncode != 0: + raise RuntimeError( + f"Failed to compute expected values (exit {proc.returncode}):\n" + f"{proc.stdout}\n{proc.stderr}" + ) + return json.loads(output.read_text(encoding="utf-8")) diff --git a/source/lmp/tests/test_lammps_dpa4_chg_spin_deepspin_pt2.py b/source/lmp/tests/test_lammps_dpa4_chg_spin_deepspin_pt2.py index 775725c703..94db6f11a6 100644 --- a/source/lmp/tests/test_lammps_dpa4_chg_spin_deepspin_pt2.py +++ b/source/lmp/tests/test_lammps_dpa4_chg_spin_deepspin_pt2.py @@ -33,10 +33,7 @@ ``test_lammps_dpa4_zbl_pt2.py`` applies to its generator geometry). """ -import json import os -import subprocess as sp -import sys import textwrap from pathlib import ( Path, @@ -47,6 +44,9 @@ from lammps import ( PyLammps, ) +from python_reference import ( + run_python_reference, +) from write_lmp_data import ( write_lmp_data_spin, ) @@ -141,7 +141,6 @@ def _compute_expected() -> None: # from SHARED_LIB_DIR, which the build-test env does not populate. infer_dir = str(pb_file.resolve().parent) script = textwrap.dedent(f"""\ - import json import sys import numpy as np @@ -175,13 +174,9 @@ def _compute_expected() -> None: "f": np.asarray(f[0]).tolist(), "fm": np.asarray(fm[0]).tolist(), }} - print(json.dumps(out)) + result = out """) - 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}") - # Warp and other backends may log to stdout before the final JSON payload. - result = json.loads(proc.stdout.strip().splitlines()[-1]) + result = run_python_reference(script) # Raw DeepEval force_mag (dE/dspin), scaled by LAMMPS's own # spin_norm / hbar unit convention (see ``_HBAR_METAL`` above) before diff --git a/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py b/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py index 32f20257df..a27729f256 100644 --- a/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py +++ b/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py @@ -43,7 +43,6 @@ """ import importlib.util -import json import os import shutil import signal @@ -60,6 +59,9 @@ from lammps import ( PyLammps, ) +from python_reference import ( + run_python_reference, +) from write_lmp_data import ( write_lmp_data, ) @@ -164,7 +166,7 @@ def _compute_expected() -> None: # 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 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" @@ -178,18 +180,14 @@ def _compute_expected() -> None: f" {atype!r},\n" " atomic=True,\n" ")\n" - "print(json.dumps({\n" + "result = {\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" + "}\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}") - # Warp and other backends may log to stdout before the final JSON payload. - result = json.loads(proc.stdout.strip().splitlines()[-1]) + result = run_python_reference(script) expected_e = result["e"] expected_ae = np.array(result["ae"]) diff --git a/source/lmp/tests/test_lammps_model_devi_pt2.py b/source/lmp/tests/test_lammps_model_devi_pt2.py index bb8649ebdb..841616b1a6 100644 --- a/source/lmp/tests/test_lammps_model_devi_pt2.py +++ b/source/lmp/tests/test_lammps_model_devi_pt2.py @@ -70,13 +70,14 @@ def _compute_expected(): """ global expected_e, expected_f, expected_v, expected_f2, expected_v2 - import json - import subprocess - import sys import textwrap + from python_reference import ( + run_python_reference, + ) + script = textwrap.dedent(f"""\ - import json, numpy as np + import numpy as np from deepmd.infer import DeepPot coord = {coord.tolist()!r} @@ -93,16 +94,9 @@ def _compute_expected(): atype, fparam=np.array(fp), aparam=np.array(ap), atomic=True, ) results.append({{"e": e[0, 0], "f": f[0].tolist(), "av": av[0].tolist()}}) - print(json.dumps(results)) + result = results """) - proc = subprocess.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}") - results = json.loads(proc.stdout.strip()) + results = run_python_reference(script) expected_e = results[0]["e"] expected_f = np.array(results[0]["f"]) # DeepPot returns virial; LAMMPS centroid/stress/atom returns stress = -virial diff --git a/source/lmp/tests/test_python_reference.py b/source/lmp/tests/test_python_reference.py new file mode 100644 index 0000000000..cab3c3b5be --- /dev/null +++ b/source/lmp/tests/test_python_reference.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""The reference result channel is independent of native initialization logs.""" + +import pytest +from python_reference import ( + run_python_reference, +) + + +@pytest.mark.parametrize("result", [{"energy": [1.0, -2.0]}, [1.0, {"force": 0.0}]]) +def test_reference_ignores_diagnostic_output(result: object) -> None: + script = ( + "import os\n" + "print('Warp initialized')\n" + "os.write(1, b'native stdout diagnostic\\n')\n" + "os.write(2, b'native stderr diagnostic\\n')\n" + f"result = {result!r}\n" + ) + assert run_python_reference(script) == result + + +def test_reference_failure_preserves_diagnostics() -> None: + with pytest.raises(RuntimeError, match="reference failed") as error: + run_python_reference( + "print('initialization diagnostic')\n" + "raise RuntimeError('reference failed')\n" + ) + assert "initialization diagnostic" in str(error.value) diff --git a/source/op/pt/dpa4c/graph_compress.cu b/source/op/pt/dpa4c/graph_compress.cu index 381a75043c..9a3e05a33a 100644 --- a/source/op/pt/dpa4c/graph_compress.cu +++ b/source/op/pt/dpa4c/graph_compress.cu @@ -114,6 +114,8 @@ struct Payload { torch::Tensor spin; torch::Tensor spin_pair; torch::Tensor spin_type; + // Per-element length scale in Å with one entry per type-table row. + torch::Tensor contact_radius; bool canonical; int64_t lmax; double table_stride; @@ -121,6 +123,10 @@ struct Payload { double rcut; double eps; double degree_floor; + // Bounds of the inner bridging switch as fractions of a pair's own length + // scale; equal fractions mean no window. + double f_inner; + double f_outer; int64_t node_begin = 0; }; @@ -158,7 +164,7 @@ Arguments build_arguments(const Payload& payload, &payload.pair_mixing, &payload.type_embedding, &payload.readout_matrices, &payload.coupling_meta, &payload.coupling_entry, &payload.coupling_value, &payload.output_mean, - &payload.output_inv_std}) { + &payload.output_inv_std, &payload.contact_radius}) { TORCH_CHECK(tensor->is_cuda() && tensor->device() == device, "dpa4c_graph_compress: every tensor input must be a CUDA " "tensor on the device of edge_vec"); @@ -185,8 +191,8 @@ Arguments build_arguments(const Payload& payload, for (const torch::Tensor* tensor : {&payload.table, &payload.pair_film, &payload.pair_mixing, &payload.type_embedding, &payload.readout_matrices, - &payload.coupling_value, &payload.output_mean, - &payload.output_inv_std}) { + &payload.coupling_value, &payload.output_mean, &payload.output_inv_std, + &payload.contact_radius}) { TORCH_CHECK(tensor->scalar_type() == torch::kFloat32, "dpa4c_graph_compress: tables and weights must be fp32"); } @@ -194,6 +200,10 @@ Arguments build_arguments(const Payload& payload, "dpa4c_graph_compress: lmax must be 2, 3, or 4"); TORCH_CHECK(payload.table_max >= payload.rcut, "dpa4c_graph_compress: table_max must cover rcut"); + TORCH_CHECK(payload.f_inner >= 0.0 && payload.f_outer >= payload.f_inner, + "dpa4c_graph_compress: the bridging switch needs 0 <= f_inner <= " + "f_outer, got ", + payload.f_inner, " and ", payload.f_outer); // The shared mode cache is sized from a compile-time maximum and the split // spline row assumes an even table width, so the rank set is closed. TORCH_CHECK(radial_modes == 0 || radial_modes == 2 || radial_modes == 4 || @@ -234,6 +244,10 @@ Arguments build_arguments(const Payload& payload, TORCH_CHECK(payload.output_mean.numel() == widths.output_width && payload.output_inv_std.numel() == widths.output_width, "dpa4c_graph_compress: invalid output calibration shape"); + TORCH_CHECK(payload.contact_radius.sizes() == + torch::IntArrayRef({static_cast(type_count)}), + "dpa4c_graph_compress: contact_radius must have shape " + "(type_count,)"); // A run covers ``node_count`` destination rows starting at ``node_begin``; // the type table is system-wide because neighbor lookups index it with // absolute source indices. Whole-system entry points additionally require @@ -295,6 +309,13 @@ Arguments build_arguments(const Payload& payload, arguments.rcut = static_cast(payload.rcut); arguments.eps = static_cast(payload.eps); arguments.degree_floor = static_cast(payload.degree_floor); + if (payload.f_outer > payload.f_inner) { + arguments.f_inner = static_cast(payload.f_inner); + arguments.f_outer = static_cast(payload.f_outer); + arguments.inverse_f_width = + static_cast(1.0 / (payload.f_outer - payload.f_inner)); + arguments.bridged = true; + } arguments.canonical = payload.canonical; arguments.index_kind = index_kind_of(payload.edge_index); arguments.edge_index = payload.edge_index.data_ptr(); @@ -319,9 +340,39 @@ Arguments build_arguments(const Payload& payload, arguments.coupling_value = payload.coupling_value.data_ptr(); arguments.output_mean = payload.output_mean.data_ptr(); arguments.output_inv_std = payload.output_inv_std.data_ptr(); + arguments.contact_radius = payload.contact_radius.data_ptr(); return arguments; } +/// Attach the analytical pair potential to one backward launch. +/// +/// \param arguments Launch bundle of the run. +/// \param pair_table Series constants of every ordered type pair with shape +/// `((T + 1)^2, 8)`, fp32. +/// \param device Device of the edge vectors, which the table must share. +/// \param pair_seed Energy cotangent of the first node of the run, fp64. +/// \param pair_energy Energy of the first node of the run, fp64; the scan +/// adds the pair energy onto it. +void attach_pair_term(Arguments& arguments, + const torch::Tensor& pair_table, + const torch::Device& device, + const double* pair_seed, + double* pair_energy) { + const long pairs = + static_cast(arguments.type_count) * arguments.type_count; + TORCH_CHECK(pair_table.is_cuda() && pair_table.device() == device && + pair_table.is_contiguous() && + pair_table.scalar_type() == torch::kFloat32 && + pair_table.sizes() == torch::IntArrayRef({pairs, 8}), + "dpa4c_graph_compress: pair_table must be a contiguous fp32 " + "CUDA tensor on the device of edge_vec with shape " + "(type_count^2, 8)"); + arguments.pair_table = pair_table.data_ptr(); + arguments.pair_seed = pair_seed; + arguments.pair_energy = pair_energy; + arguments.bridged = true; +} + } // namespace std::tuple dpa4c_graph_compress( @@ -344,13 +395,16 @@ std::tuple dpa4c_graph_compress( torch::Tensor spin, torch::Tensor spin_pair, torch::Tensor spin_type, + torch::Tensor contact_radius, bool canonical, int64_t lmax, double table_stride, double table_max, double rcut, double eps, - double degree_floor) { + double degree_floor, + double f_inner, + double f_outer) { const Payload payload{edge_index, edge_mask, destination_order, @@ -369,13 +423,16 @@ std::tuple dpa4c_graph_compress( spin, spin_pair, spin_type, + contact_radius, canonical, lmax, table_stride, table_max, rcut, eps, - degree_floor}; + degree_floor, + f_inner, + f_outer}; const long node_count = destination_row_ptr.numel() - 1; // The destination row pointer defines the node axis; the type table must // describe exactly that axis, or the two disagree on how many nodes exist. @@ -403,7 +460,7 @@ std::tuple dpa4c_graph_compress( return {descriptor, state}; } -std::tuple +std::tuple dpa4c_graph_compress_backward_impl(torch::Tensor descriptor_gradient, torch::Tensor state, torch::Tensor edge_vec, @@ -425,6 +482,7 @@ dpa4c_graph_compress_backward_impl(torch::Tensor descriptor_gradient, torch::Tensor spin, torch::Tensor spin_pair, torch::Tensor spin_type, + torch::Tensor contact_radius, bool canonical, int64_t lmax, double table_stride, @@ -432,6 +490,10 @@ dpa4c_graph_compress_backward_impl(torch::Tensor descriptor_gradient, double rcut, double eps, double degree_floor, + double f_inner, + double f_outer, + torch::Tensor pair_table, + torch::Tensor pair_seed, bool reuse_state) { const Payload payload{edge_index, edge_mask, @@ -451,13 +513,16 @@ dpa4c_graph_compress_backward_impl(torch::Tensor descriptor_gradient, spin, spin_pair, spin_type, + contact_radius, canonical, lmax, table_stride, table_max, rcut, eps, - degree_floor}; + degree_floor, + f_inner, + f_outer}; const bool has_spin = spin.dim() == 2; const long node_count = destination_row_ptr.numel() - 1; // The destination row pointer defines the node axis; the type table must @@ -484,20 +549,27 @@ dpa4c_graph_compress_backward_impl(torch::Tensor descriptor_gradient, descriptor_gradient.numel() == node_count * widths.output_width, "dpa4c_graph_compress_backward: invalid descriptor gradient"); auto float_options = edge_vec.options().dtype(torch::kFloat32); - // Every absent output receives its own allocation. The schema declares three + // Every absent output receives its own allocation. The schema declares four // unannotated results, so returning one empty tensor in two slots would // introduce an alias the schema does not describe, which is undefined under - // functionalization for all three inputs rather than only for spin. + // functionalization whether or not the spin branch is active. const auto absent = [&float_options] { return torch::empty({0}, float_options); }; + // The pair energy accumulates onto zeros here; a model without an analytical + // pair potential passes an empty table and receives an empty energy. + const bool has_pair = pair_table.numel() != 0; + auto pair_energy = + has_pair ? torch::zeros({node_count, 1}, + edge_vec.options().dtype(torch::kFloat64)) + : torch::empty({0}, edge_vec.options().dtype(torch::kFloat64)); if (node_count == 0) { if (has_spin) { return {torch::zeros_like(edge_vec), torch::empty({node_count, 3}, float_options), - torch::empty(edge_vec.sizes(), float_options)}; + torch::empty(edge_vec.sizes(), float_options), pair_energy}; } - return {torch::zeros_like(edge_vec), absent(), absent()}; + return {torch::zeros_like(edge_vec), absent(), absent(), pair_energy}; } auto descriptor_gradient_float = descriptor_gradient.to(torch::kFloat32).contiguous(); @@ -523,12 +595,25 @@ dpa4c_graph_compress_backward_impl(torch::Tensor descriptor_gradient, arguments.spin_gradient = spin_gradient.data_ptr(); arguments.edge_spin_gradient = edge_spin_gradient.data_ptr(); } + auto pair_seed_c = pair_seed.contiguous(); + if (has_pair) { + TORCH_CHECK(pair_seed_c.is_cuda() && + pair_seed_c.device() == edge_vec.device() && + pair_seed_c.numel() == node_count && + pair_seed_c.scalar_type() == torch::kFloat64, + "dpa4c_graph_compress_backward: pair_seed must be an fp64 " + "CUDA tensor on the device of edge_vec with one entry per " + "node"); + attach_pair_term(arguments, pair_table, edge_vec.device(), + pair_seed_c.data_ptr(), + pair_energy.data_ptr()); + } dispatch(channels, true, arguments, at::cuda::getCurrentCUDAStream()); return {edge_gradient.to(edge_vec.scalar_type()), spin_gradient, - edge_spin_gradient}; + edge_spin_gradient, pair_energy}; } -std::tuple +std::tuple dpa4c_graph_compress_backward(torch::Tensor descriptor_gradient, torch::Tensor state, torch::Tensor edge_vec, @@ -550,20 +635,26 @@ dpa4c_graph_compress_backward(torch::Tensor descriptor_gradient, torch::Tensor spin, torch::Tensor spin_pair, torch::Tensor spin_type, + torch::Tensor contact_radius, bool canonical, int64_t lmax, double table_stride, double table_max, double rcut, double eps, - double degree_floor) { + double degree_floor, + double f_inner, + double f_outer, + torch::Tensor pair_table, + torch::Tensor pair_seed) { return dpa4c_graph_compress_backward_impl( descriptor_gradient, state, edge_vec, edge_index, edge_mask, destination_order, destination_row_ptr, atype, table, pair_film, pair_mixing, type_embedding, readout_matrices, coupling_meta, coupling_entry, coupling_value, output_mean, output_inv_std, spin, - spin_pair, spin_type, canonical, lmax, table_stride, table_max, rcut, eps, - degree_floor, false); + spin_pair, spin_type, contact_radius, canonical, lmax, table_stride, + table_max, rcut, eps, degree_floor, f_inner, f_outer, pair_table, + pair_seed, false); } std::tuple dpa4c_canonical_compress( @@ -584,12 +675,15 @@ std::tuple dpa4c_canonical_compress( torch::Tensor spin, torch::Tensor spin_pair, torch::Tensor spin_type, + torch::Tensor contact_radius, int64_t lmax, double table_stride, double table_max, double rcut, double eps, - double degree_floor) { + double degree_floor, + double f_inner, + double f_outer) { TORCH_CHECK(source.dim() == 1 && source.numel() == edge_vec.size(0), "dpa4c_canonical_compress: source and edge_vec must share the " "edge axis"); @@ -599,8 +693,8 @@ std::tuple dpa4c_canonical_compress( edge_vec, source, edge_mask, destination_order, destination_row_ptr, atype, table, pair_film, pair_mixing, type_embedding, readout_matrices, coupling_meta, coupling_entry, coupling_value, output_mean, - output_inv_std, spin, spin_pair, spin_type, true, lmax, table_stride, - table_max, rcut, eps, degree_floor); + output_inv_std, spin, spin_pair, spin_type, contact_radius, true, lmax, + table_stride, table_max, rcut, eps, degree_floor, f_inner, f_outer); } std::tuple @@ -623,25 +717,34 @@ dpa4c_canonical_compress_backward_common(torch::Tensor descriptor_gradient, torch::Tensor spin, torch::Tensor spin_pair, torch::Tensor spin_type, + torch::Tensor contact_radius, int64_t lmax, double table_stride, double table_max, double rcut, double eps, double degree_floor, + double f_inner, + double f_outer, bool reuse_state) { TORCH_CHECK(source.dim() == 1 && source.numel() == edge_vec.size(0), "dpa4c_canonical_compress_backward: source and edge_vec must " "share the edge axis"); auto edge_mask = torch::empty({0}, edge_vec.options().dtype(torch::kBool)); auto destination_order = torch::empty({0}, source.options()); - return dpa4c_graph_compress_backward_impl( - descriptor_gradient, state, edge_vec, source, edge_mask, - destination_order, destination_row_ptr, atype, table, pair_film, - pair_mixing, type_embedding, readout_matrices, coupling_meta, - coupling_entry, coupling_value, output_mean, output_inv_std, spin, - spin_pair, spin_type, true, lmax, table_stride, table_max, rcut, eps, - degree_floor, reuse_state); + // The compact descriptor operators carry no analytical pair potential; the + // fused energy-gradient operator evaluates it on this ABI. + auto no_pair = torch::empty({0}, edge_vec.options().dtype(torch::kFloat32)); + auto [edge_gradient, spin_gradient, edge_spin_gradient, pair_energy] = + dpa4c_graph_compress_backward_impl( + descriptor_gradient, state, edge_vec, source, edge_mask, + destination_order, destination_row_ptr, atype, table, pair_film, + pair_mixing, type_embedding, readout_matrices, coupling_meta, + coupling_entry, coupling_value, output_mean, output_inv_std, spin, + spin_pair, spin_type, contact_radius, true, lmax, table_stride, + table_max, rcut, eps, degree_floor, f_inner, f_outer, no_pair, + no_pair, reuse_state); + return {edge_gradient, spin_gradient, edge_spin_gradient}; } std::tuple @@ -664,18 +767,22 @@ dpa4c_canonical_compress_backward(torch::Tensor descriptor_gradient, torch::Tensor spin, torch::Tensor spin_pair, torch::Tensor spin_type, + torch::Tensor contact_radius, int64_t lmax, double table_stride, double table_max, double rcut, double eps, - double degree_floor) { + double degree_floor, + double f_inner, + double f_outer) { return dpa4c_canonical_compress_backward_common( descriptor_gradient, state, edge_vec, source, destination_row_ptr, atype, table, pair_film, pair_mixing, type_embedding, readout_matrices, coupling_meta, coupling_entry, coupling_value, output_mean, - output_inv_std, spin, spin_pair, spin_type, lmax, table_stride, table_max, - rcut, eps, degree_floor, false); + output_inv_std, spin, spin_pair, spin_type, contact_radius, lmax, + table_stride, table_max, rcut, eps, degree_floor, f_inner, f_outer, + false); } std::tuple @@ -698,18 +805,21 @@ dpa4c_canonical_compress_backward_inplace(torch::Tensor descriptor_gradient, torch::Tensor spin, torch::Tensor spin_pair, torch::Tensor spin_type, + torch::Tensor contact_radius, int64_t lmax, double table_stride, double table_max, double rcut, double eps, - double degree_floor) { + double degree_floor, + double f_inner, + double f_outer) { return dpa4c_canonical_compress_backward_common( descriptor_gradient, state, edge_vec, source, destination_row_ptr, atype, table, pair_film, pair_mixing, type_embedding, readout_matrices, coupling_meta, coupling_entry, coupling_value, output_mean, - output_inv_std, spin, spin_pair, spin_type, lmax, table_stride, table_max, - rcut, eps, degree_floor, true); + output_inv_std, spin, spin_pair, spin_type, contact_radius, lmax, + table_stride, table_max, rcut, eps, degree_floor, f_inner, f_outer, true); } // Energy and edge cotangent of one compressed inference step, evaluated over @@ -743,12 +853,15 @@ dpa4c_canonical_compress_energy_gradient(torch::Tensor edge_vec, torch::Tensor spin, torch::Tensor spin_pair, torch::Tensor spin_type, + torch::Tensor contact_radius, int64_t lmax, double table_stride, double table_max, double rcut, double eps, double degree_floor, + double f_inner, + double f_outer, std::vector ws, std::vector bs, std::vector resnets, @@ -757,7 +870,8 @@ dpa4c_canonical_compress_energy_gradient(torch::Tensor edge_vec, torch::Tensor bias_atom_e, int64_t act, torch::Tensor seed, - int64_t tile) { + int64_t tile, + torch::Tensor pair_table) { TORCH_CHECK(edge_vec.is_cuda(), "dpa4c_canonical_compress_energy_gradient: edge_vec must be a " "CUDA tensor"); @@ -825,19 +939,35 @@ dpa4c_canonical_compress_energy_gradient(torch::Tensor edge_vec, for (long begin = 0; begin < node_count; begin += run) { const long count = std::min(run, node_count - begin); const Payload payload{ - source, empty_mask, - empty_index, destination_row_ptr.slice(0, begin, begin + count + 1), - atype, table, - pair_film, pair_mixing, - type_embedding, readout_matrices, - coupling_meta, coupling_entry, - coupling_value, output_mean, - output_inv_std, spin, - spin_pair, spin_type, - true, lmax, - table_stride, table_max, - rcut, eps, - degree_floor, begin}; + source, + empty_mask, + empty_index, + destination_row_ptr.slice(0, begin, begin + count + 1), + atype, + table, + pair_film, + pair_mixing, + type_embedding, + readout_matrices, + coupling_meta, + coupling_entry, + coupling_value, + output_mean, + output_inv_std, + spin, + spin_pair, + spin_type, + contact_radius, + true, + lmax, + table_stride, + table_max, + rcut, + eps, + degree_floor, + f_inner, + f_outer, + begin}; Arguments arguments = build_arguments(payload, edge_vec_float, channels, widths); arguments.descriptor = descriptor.data_ptr(); @@ -867,6 +997,13 @@ dpa4c_canonical_compress_energy_gradient(torch::Tensor edge_vec, // Only the final run reaches the reserved edge slots; its row pointer ends // at the last physical edge, which is exactly where the padding begins. arguments.clear_padding = begin + count == node_count; + if (pair_table.numel() != 0) { + // The fitting forward has assigned the energy rows of this run, and the + // seed is the cotangent the fitting backward just consumed. + attach_pair_term(arguments, pair_table, edge_vec.device(), + seed_c.data_ptr() + begin, + energy.data_ptr() + begin); + } dispatch(channels, true, arguments, stream); } return {energy, edge_gradient.to(edge_vec.scalar_type()), spin_gradient, diff --git a/source/op/pt/dpa4c/graph_compress.cuh b/source/op/pt/dpa4c/graph_compress.cuh index 5ad8550238..6d211836d2 100644 --- a/source/op/pt/dpa4c/graph_compress.cuh +++ b/source/op/pt/dpa4c/graph_compress.cuh @@ -215,6 +215,62 @@ __device__ __forceinline__ void accumulate_modes_with_derivative( } } +// The edge envelope is the C³ cutoff factor times the inner bridging switch, +// and the radial table of a bridged model is read at a clamped length rather +// than the true one. The window of a pair is two dimensionless fractions of +// that pair's own length scale `s`, the sum of the covalent radii of its two +// elements, so both are functions of the reduced distance `u = r / s` through +// the window parameter `t = clamp((u - f_inner) / (f_outer - f_inner), 0, 1)`: +// +// w(u) = h(t), h(t) = 35 t^4 - 84 t^5 + 70 t^6 - 20 t^7, +// h'(t) = 140 t^3 (1 - t)^3, +// ũ(u) = f_mid + (f_outer - f_inner) S(t), +// S(t) = 7 t^5 - 14 t^6 + 10 t^7 - 2.5 t^8, +// f_mid = (f_inner + f_outer) / 2, +// +// the septic smootherstep of the portable descriptor and its antiderivative, +// evaluated in the same Horner forms so that all three agree term by term. +// The switch vanishes exactly at and below the inner fraction, which removes +// the edge from every moment, and equals one exactly at and beyond the outer +// fraction. Since `S' = h`, `S(0) = 0` and `S(1) = 1/2`, the clamped length +// `r̃ = s ũ` is frozen at the window midpoint at and below the inner fraction, +// equals `r` at and beyond the outer fraction, and moves with `r` at the rate +// `dr̃/dr = w` in between. The kernels therefore evaluate the window only for +// an edge inside it, which the comparison `r < f_outer * s` decides without a +// division: in a simulation nearly every edge lies outside, so the edges of a +// warp almost always skip it together and the common path costs one +// multiply-compare. +// +// `bridging_switch_derivative` returns dh/du; a caller that needs the radial +// derivative divides it by the pair's length scale once. + +__device__ __forceinline__ float window_parameter(float reduced, + float inner, + float inverse_width) { + return fminf(fmaxf((reduced - inner) * inverse_width, 0.0f), 1.0f); +} + +__device__ __forceinline__ float bridging_switch(float t) { + const float t2 = t * t; + return t2 * t2 * (35.0f + t * (-84.0f + t * (70.0f - 20.0f * t))); +} + +__device__ __forceinline__ float bridging_switch_derivative( + float t, float inverse_width) { + const float bump = t * (1.0f - t); + return 140.0f * bump * bump * bump * inverse_width; +} + +__device__ __forceinline__ float clamped_length(float t, + float contact, + float inner, + float outer) { + const float t2 = t * t; + const float rise = + t2 * t2 * t * (7.0f + t * (-14.0f + t * (10.0f - 2.5f * t))); + return contact * fmaf(outer - inner, rise, 0.5f * (inner + outer)); +} + __device__ __forceinline__ float c3_envelope(float radius, float rcut) { const float u = fminf(fmaxf(__fdividef(rcut - radius, rcut), 0.0f), 1.0f); const float x = 1.0f - u; @@ -240,31 +296,85 @@ __device__ __forceinline__ float c3_envelope_derivative(float radius, return __fdividef(-4.0f * u2 * u * series + u2 * u2 * derivative, rcut); } +// === Analytical pair potential === +// +// One ordered type pair owns the eight constants [A_0..A_3, c_0..c_3] of the +// screened Coulomb series (the ZBL form) +// +// V(r) = (1/r) sum_k A_k exp(-c_k r), +// dV/dr = -(1/r) sum_k A_k (c_k + 1/r) exp(-c_k r). +// +// Energy and slope share their exponentials, so both come from one +// evaluation. The `Width` lanes of an edge group split the four terms between +// them, lane `l` taking the terms `k = l (mod Width)`: the partial slopes are +// merged by the reduction that already merges the lanes of an edge and the +// partial energies by one reduction over the warp after the scan, no lane +// evaluates more than two exponentials per edge, and the split is resolved at +// compile time. + +struct PairTerm { + float energy; + float slope; +}; + +template +__device__ __forceinline__ PairTerm +pair_term(const float* row, int lane, float radius, float inverse_radius) { + PairTerm term = {0.0f, 0.0f}; +#pragma unroll + for (int first = 0; first < 4; first += Width) { + const int k = first + lane; + if (Width <= 4 || k < 4) { + const float rate = __ldg(row + 4 + k); + const float decayed = + __ldg(row + k) * expf(-rate * radius) * inverse_radius; + term.energy += decayed; + term.slope = fmaf(-decayed, rate + inverse_radius, term.slope); + } + } + return term; +} + // === Edge geometry and angular basis === struct EdgeGeometry { float ux; float uy; float uz; + // True regularized length, which the direction, the analytical pair + // potential and the bridging window read. float radius; float inverse_radius; + // Length the radial table reads: the clamped length inside a bridging + // window, the true one everywhere else. + float seen; float envelope; int source_type; // Retained for the native spin branch, which gathers the source moment. long source; }; -template -__device__ __forceinline__ EdgeGeometry load_geometry(long edge, - float rcut, - float eps, - const float* edge_vec, - const index_t* edge_index, - const long* atype) { +// The length scale of a pair is the sum of the two per-element radii. The +// centre contributes a block constant, so only the source radius is indexed +// per edge; `EdgeGeometry` keeps the source type, from which the backward +// recovers the same sum, given that constant, at the cost of one further load +// and one add. +template +__device__ __forceinline__ EdgeGeometry +load_geometry(long edge, + long source, + int source_type, + float rcut, + float f_inner, + float f_outer, + float inverse_f_width, + float center_radius, + const float* contact_radius, + float eps, + const float* edge_vec) { EdgeGeometry geometry; - const long source = static_cast(edge_index[edge]); geometry.source = source; - geometry.source_type = static_cast(atype[source]); + geometry.source_type = source_type; const float x = edge_vec[edge * 3 + 0]; const float y = edge_vec[edge * 3 + 1]; const float z = edge_vec[edge * 3 + 2]; @@ -274,10 +384,61 @@ __device__ __forceinline__ EdgeGeometry load_geometry(long edge, geometry.ux = x * geometry.inverse_radius; geometry.uy = y * geometry.inverse_radius; geometry.uz = z * geometry.inverse_radius; + geometry.seen = geometry.radius; geometry.envelope = c3_envelope(geometry.radius, rcut); + if constexpr (Bridged) { + const float contact = + center_radius + __ldg(contact_radius + geometry.source_type); + if (geometry.radius < f_outer * contact) { + const float t = window_parameter(__fdividef(geometry.radius, contact), + f_inner, inverse_f_width); + geometry.seen = clamped_length(t, contact, f_inner, f_outer); + geometry.envelope = c3_envelope(geometry.seen, rcut) * bridging_switch(t); + } + } return geometry; } +// Closes the radial chain rule at the tail of the backward scan. The table +// cotangent differentiates the length the table reads and the envelope +// cotangent the envelope; both become derivatives with respect to the true +// length. In a plain model, and outside the window of a bridged one, the +// table reads the true length and the envelope is the cutoff factor alone. +// Inside the window the table reads the clamped length, whose slope is the +// switch `w`, and the envelope is `c3(r̃) w`, so +// +// d/dr = table_gradient * w + envelope_gradient * (c3'(r̃) w^2 + c3(r̃) dw/dr), +// +// with `dw/dr = h'(t) / ((f_outer - f_inner) s)`. The window is resolved again +// here rather than carried through the channel scan, which keeps its state +// out of the registers the scan is short of. +template +__device__ __forceinline__ float radial_chain(const EdgeGeometry& geometry, + const Arguments& args, + float center_radius, + float table_gradient, + float envelope_gradient) { + if constexpr (Bridged) { + const float contact = + center_radius + __ldg(args.contact_radius + geometry.source_type); + if (geometry.radius < args.f_outer * contact) { + const float t = window_parameter(__fdividef(geometry.radius, contact), + args.f_inner, args.inverse_f_width); + const float seen = clamped_length(t, contact, args.f_inner, args.f_outer); + const float weight = bridging_switch(t); + const float slope = fmaf( + c3_envelope_derivative(seen, args.rcut) * weight, weight, + c3_envelope(seen, args.rcut) * + __fdividef(bridging_switch_derivative(t, args.inverse_f_width), + contact)); + return fmaf(envelope_gradient, slope, table_gradient * weight); + } + } + return fmaf(envelope_gradient, + c3_envelope_derivative(geometry.radius, args.rcut), + table_gradient); +} + // The Cartesian harmonics are evaluated on the unit direction, so the squared // norm that makes each degree traceless is exactly one. Two polynomials that // agree on the unit sphere differ by a multiple of `|u|^2 - 1`, whose gradient diff --git a/source/op/pt/dpa4c/graph_compress_cpu.cc b/source/op/pt/dpa4c/graph_compress_cpu.cc index 20e98bb2ec..658aed0759 100644 --- a/source/op/pt/dpa4c/graph_compress_cpu.cc +++ b/source/op/pt/dpa4c/graph_compress_cpu.cc @@ -203,14 +203,14 @@ int isa_block() { } /// Resolve the kernels of the running CPU. -Kernels resolve_kernels(int lmax, bool has_modes) { +Kernels resolve_kernels(int lmax, bool has_modes, bool bridged) { switch (deepmd_cpu::host_isa()) { case Isa::kAvx512: - return avx512::kernels(lmax, has_modes); + return avx512::kernels(lmax, has_modes, bridged); case Isa::kAvx2: - return avx2::kernels(lmax, has_modes); + return avx2::kernels(lmax, has_modes, bridged); default: - return scalar::kernels(lmax, has_modes); + return scalar::kernels(lmax, has_modes, bridged); } } @@ -333,11 +333,14 @@ Arguments build_arguments(const Inputs& inputs, const torch::Tensor& coupling_value, const torch::Tensor& output_mean, const torch::Tensor& output_inv_std, + const torch::Tensor& contact_radius, double table_stride, double table_max, double rcut, double eps, - double degree_floor) { + double degree_floor, + double f_inner, + double f_outer) { Arguments arguments{}; arguments.edge_vec = inputs.edge_vec.const_data_ptr(); arguments.source = inputs.source.const_data_ptr(); @@ -364,6 +367,14 @@ Arguments build_arguments(const Inputs& inputs, : coupling_value.const_data_ptr(); arguments.output_mean = output_mean.const_data_ptr(); arguments.output_inv_std = output_inv_std.const_data_ptr(); + TORCH_CHECK( + contact_radius.is_contiguous() && + contact_radius.scalar_type() == torch::kFloat32 && + contact_radius.sizes() == + torch::IntArrayRef({static_cast(inputs.layout.type_count)}), + "dpa4c_graph_compress: contact_radius must be a contiguous fp32 " + "tensor with shape (type_count,)"); + arguments.contact_radius = contact_radius.const_data_ptr(); arguments.node_count = inputs.atype.size(0); arguments.edge_count = inputs.edge_vec.size(0); arguments.coupling_count = static_cast(coupling_meta.numel() / 8); @@ -372,6 +383,15 @@ Arguments build_arguments(const Inputs& inputs, arguments.rcut = static_cast(rcut); arguments.eps = static_cast(eps); arguments.degree_floor = static_cast(degree_floor); + TORCH_CHECK(f_inner >= 0.0 && f_outer >= f_inner, + "dpa4c_graph_compress: the bridging switch needs 0 <= " + "f_inner <= f_outer, got ", + f_inner, " and ", f_outer); + const bool windowed = f_outer > f_inner; + arguments.f_inner = windowed ? static_cast(f_inner) : 0.0f; + arguments.f_outer = windowed ? static_cast(f_outer) : 0.0f; + arguments.inverse_f_width = + windowed ? static_cast(1.0 / (f_outer - f_inner)) : 0.0f; return arguments; } @@ -411,13 +431,16 @@ std::tuple forward( torch::Tensor spin, torch::Tensor spin_pair, torch::Tensor spin_type, + torch::Tensor contact_radius, bool canonical, int64_t lmax, double table_stride, double table_max, double rcut, double eps, - double degree_floor) { + double degree_floor, + double f_inner, + double f_outer) { TORCH_CHECK(spin.dim() != 2, "dpa4c_graph_compress: the CPU kernel has no native spin " "branch; evaluate a spin-conditioned descriptor eagerly"); @@ -430,20 +453,22 @@ std::tuple forward( torch::empty({inputs.atype.size(0), layout.output_width}, options); auto state = torch::empty({inputs.atype.size(0), layout.moment_width + 2}, options); - Arguments arguments = - build_arguments(inputs, type_embedding.contiguous(), - readout_matrices.contiguous(), coupling_meta.contiguous(), - coupling_entry.contiguous(), coupling_value.contiguous(), - output_mean.contiguous(), output_inv_std.contiguous(), - table_stride, table_max, rcut, eps, degree_floor); + Arguments arguments = build_arguments( + inputs, type_embedding.contiguous(), readout_matrices.contiguous(), + coupling_meta.contiguous(), coupling_entry.contiguous(), + coupling_value.contiguous(), output_mean.contiguous(), + output_inv_std.contiguous(), contact_radius, table_stride, table_max, + rcut, eps, degree_floor, f_inner, f_outer); arguments.descriptor = descriptor.data_ptr(); arguments.state = state.data_ptr(); - run_scan(resolve_kernels(static_cast(lmax), layout.modes > 0).forward, + run_scan(resolve_kernels(static_cast(lmax), layout.modes > 0, + f_outer > f_inner) + .forward, arguments, layout); return {descriptor, state}; } -std::tuple backward( +std::tuple backward( torch::Tensor descriptor_gradient, torch::Tensor state, torch::Tensor edge_vec, @@ -465,13 +490,18 @@ std::tuple backward( torch::Tensor spin, torch::Tensor spin_pair, torch::Tensor spin_type, + torch::Tensor contact_radius, bool canonical, int64_t lmax, double table_stride, double table_max, double rcut, double eps, - double degree_floor) { + double degree_floor, + double f_inner, + double f_outer, + torch::Tensor pair_table, + torch::Tensor pair_seed) { TORCH_CHECK(spin.dim() != 2, "dpa4c_graph_compress_backward: the CPU kernel has no native " "spin branch"); @@ -482,12 +512,12 @@ std::tuple backward( auto options = edge_vec.options().dtype(torch::kFloat32); auto edge_gradient = torch::empty({inputs.edge_vec.size(0), 3}, options); auto absent = torch::empty({0}, options); - Arguments arguments = - build_arguments(inputs, type_embedding.contiguous(), - readout_matrices.contiguous(), coupling_meta.contiguous(), - coupling_entry.contiguous(), coupling_value.contiguous(), - output_mean.contiguous(), output_inv_std.contiguous(), - table_stride, table_max, rcut, eps, degree_floor); + Arguments arguments = build_arguments( + inputs, type_embedding.contiguous(), readout_matrices.contiguous(), + coupling_meta.contiguous(), coupling_entry.contiguous(), + coupling_value.contiguous(), output_mean.contiguous(), + output_inv_std.contiguous(), contact_radius, table_stride, table_max, + rcut, eps, degree_floor, f_inner, f_outer); auto contiguous_gradient = descriptor_gradient.to(torch::kFloat32).contiguous(); auto contiguous_state = state.to(torch::kFloat32).contiguous(); @@ -496,6 +526,32 @@ std::tuple backward( const_cast(contiguous_state.const_data_ptr()); arguments.edge_gradient = edge_gradient.data_ptr(); + // The pair energy accumulates onto zeros here; a model without an analytical + // pair potential passes an empty table and receives an empty energy. + const bool has_pair = pair_table.numel() != 0; + const int64_t node_count = inputs.atype.size(0); + auto pair_energy = + has_pair ? torch::zeros({node_count, 1}, + edge_vec.options().dtype(torch::kFloat64)) + : torch::empty({0}, edge_vec.options().dtype(torch::kFloat64)); + auto contiguous_table = pair_table.contiguous(); + auto contiguous_seed = pair_seed.contiguous(); + if (has_pair) { + const int64_t pairs = + static_cast(layout.type_count) * layout.type_count; + TORCH_CHECK(contiguous_table.scalar_type() == torch::kFloat32 && + contiguous_table.sizes() == torch::IntArrayRef({pairs, 8}), + "dpa4c_graph_compress_backward: pair_table must be fp32 with " + "shape (type_count^2, 8)"); + TORCH_CHECK(contiguous_seed.scalar_type() == torch::kFloat64 && + contiguous_seed.numel() == node_count, + "dpa4c_graph_compress_backward: pair_seed must be fp64 with " + "one entry per node"); + arguments.pair_table = contiguous_table.const_data_ptr(); + arguments.pair_seed = contiguous_seed.const_data_ptr(); + arguments.pair_energy = pair_energy.data_ptr(); + } + // A masked edge sorts past the last destination row, so no row reaches it // and the scan never writes its slot. Both topology forms keep those slots // in the suffix of the destination permutation, which is the identity for @@ -520,9 +576,12 @@ std::tuple backward( }); } } - run_scan(resolve_kernels(static_cast(lmax), layout.modes > 0).backward, + run_scan(resolve_kernels(static_cast(lmax), layout.modes > 0, + f_outer > f_inner || has_pair) + .backward, arguments, layout); - return {edge_gradient.to(edge_vec.scalar_type()), absent, absent.clone()}; + return {edge_gradient.to(edge_vec.scalar_type()), absent, absent.clone(), + pair_energy}; } } // namespace diff --git a/source/op/pt/dpa4c/graph_compress_cpu.h b/source/op/pt/dpa4c/graph_compress_cpu.h index 66b02195b0..45708f04d1 100644 --- a/source/op/pt/dpa4c/graph_compress_cpu.h +++ b/source/op/pt/dpa4c/graph_compress_cpu.h @@ -188,6 +188,10 @@ struct Arguments { const float* coupling_value; ///< Gaunt values and probe scales. const float* output_mean; ///< `(D,)` calibration shift. const float* output_inv_std; ///< `(D,)` calibration scale. + /// Per-element length scale in Å with one entry per type-table row, + /// including the padding row. The sum of the two entries of a pair is the + /// distance the window fractions are measured against. + const float* contact_radius; float* descriptor; ///< `(N, D)` output. float* state; ///< `(N, S + 2)` saved state. @@ -204,6 +208,24 @@ struct Arguments { float rcut; float eps; float degree_floor; + + /// Inner bridging switch of the edge envelope: the two dimensionless + /// fractions of a pair's own length scale that bound its window and the + /// reciprocal of their difference. A bridged launch without a window keeps + /// the outer fraction at zero, which no edge lies below. + float f_inner; + float f_outer; + float inverse_f_width; + + /// Analytical pair potential, present together or not at all and read by + /// the backward only. `pair_table` holds, per ordered type pair, the four + /// amplitudes and the four decay rates of the screened Coulomb series + /// \f$V(r)=r^{-1}\sum_k A_k e^{-c_k r}\f$; `pair_seed` is the energy + /// cotangent of each node and `pair_energy` the fp64 node energies onto + /// which the scan adds \f$V/2\f$ per edge. + const float* pair_table; ///< `((T + 1)^2, 8)` or null. + const double* pair_seed; ///< `(N,)`. + double* pair_energy; ///< `(N,)`. }; /// Evaluate the descriptor over one contiguous node range. @@ -219,15 +241,15 @@ struct Kernels { }; namespace scalar { -Kernels kernels(int lmax, bool has_modes); +Kernels kernels(int lmax, bool has_modes, bool bridged); } // namespace scalar namespace avx2 { -Kernels kernels(int lmax, bool has_modes); +Kernels kernels(int lmax, bool has_modes, bool bridged); } // namespace avx2 namespace avx512 { -Kernels kernels(int lmax, bool has_modes); +Kernels kernels(int lmax, bool has_modes, bool bridged); } // namespace avx512 } // namespace deepmd_dpa4c_cpu diff --git a/source/op/pt/dpa4c/graph_compress_cpu_kernel.h b/source/op/pt/dpa4c/graph_compress_cpu_kernel.h index 6794380751..eb5480d585 100644 --- a/source/op/pt/dpa4c/graph_compress_cpu_kernel.h +++ b/source/op/pt/dpa4c/graph_compress_cpu_kernel.h @@ -37,6 +37,16 @@ #include "graph_compress_cpu.h" +// Every helper an edge scan calls per edge is inlined unconditionally. The +// scans are instantiated per angular degree, mode residual and bridging +// variant, and the unit-growth budget of the inliner would otherwise move +// helpers out of line, behind a call on every edge. +#if defined(_MSC_VER) +#define DPA4C_CPU_INLINE __forceinline +#else +#define DPA4C_CPU_INLINE inline __attribute__((always_inline)) +#endif + namespace deepmd_dpa4c_cpu { namespace DPA4C_CPU_ISA { @@ -68,13 +78,24 @@ constexpr int angular_components(int lmax) { // === Cutoff envelope === -/// Evaluate the exponent-five C³ envelope and its derivative. +/// Radial state of one edge: the envelope and the length the table reads, +/// each with its derivative with respect to the true length. +struct EdgeEnvelope { + float chi; ///< Envelope value. + float chi_slope; ///< d(chi)/d(radius) in 1/Å. + float seen; ///< Length the radial table reads, in Å. + float chain; ///< d(seen)/d(radius). +}; + +/// Evaluate the exponent-five C³ cutoff factor and its derivative. /// -/// \param radius Regularized distance in Å. +/// \param radius Length in Å. /// \param rcut Outer cutoff in Å. -/// \param derivative Receives d(envelope)/d(radius). -/// \return The envelope value, exactly zero at and beyond the cutoff. -inline float envelope(float radius, float rcut, float* derivative) { +/// \param derivative Receives d(cutoff)/d(radius) in 1/Å. +/// \return The cutoff factor, exactly zero at and beyond the cutoff. +DPA4C_CPU_INLINE float cutoff_factor(float radius, + float rcut, + float* derivative) { if (radius >= rcut) { *derivative = 0.0f; return 0.0f; @@ -92,6 +113,72 @@ inline float envelope(float radius, float rcut, float* derivative) { return u4 * series; } +/// Evaluate the radial state of one edge. +/// +/// A plain specialization reads the table at the true length and takes the +/// cutoff factor as the envelope. A bridged specialization resolves the +/// window of the pair, two fractions of its length scale \f$s\f$, through +/// the window parameter \f$t=\mathrm{clamp}((r/s-f_\mathrm{inner})/ +/// (f_\mathrm{outer}-f_\mathrm{inner}),0,1)\f$: the inner bridging switch +/// \f$h(t)=35t^4-84t^5+70t^6-20t^7\f$ multiplies the cutoff factor, and the +/// table reads the clamped length +/// \f$\tilde r=s\,(f_\mathrm{mid}+(f_\mathrm{outer}-f_\mathrm{inner})S(t))\f$ +/// with \f$S(t)=7t^5-14t^6+10t^7-2.5t^8\f$ and \f$f_\mathrm{mid}\f$ the window +/// midpoint. Both are the polynomials of the portable descriptor in the same +/// Horner forms. \f$S\f$ is the antiderivative of \f$h\f$ with \f$S(1)=1/2\f$, +/// so the clamped length is frozen at the midpoint at and below the inner +/// fraction, equals the true length at and beyond the outer fraction, and +/// moves with it at the rate \f$h(t)\f$ in between; the envelope is then +/// \f$c(\tilde r)\,h(t)\f$ with the cutoff factor \f$c\f$, differentiated by +/// the product and chain rules. An edge at or beyond the outer fraction takes +/// the plain path. +/// +/// \param radius Regularized distance in Å. +/// \param contact Length scale \f$s\f$ of the pair in Å. +/// \param rcut Outer cutoff in Å. +/// \param f_inner Inner fraction of the pair length scale. +/// \param f_outer Outer fraction of the pair length scale. +/// \param inverse_f_width Reciprocal difference of the two fractions. +/// \return The radial state; the envelope is exactly zero at and beyond the +/// cutoff and, when bridged, at and below the inner fraction. +template +DPA4C_CPU_INLINE EdgeEnvelope envelope(float radius, + float contact, + float rcut, + float f_inner, + float f_outer, + float inverse_f_width) { + EdgeEnvelope state; + if (!BRIDGED || radius >= f_outer * contact) { + state.chi = cutoff_factor(radius, rcut, &state.chi_slope); + state.seen = radius; + state.chain = 1.0f; + return state; + } + + const float inverse_contact = 1.0f / contact; + const float reduced = radius * inverse_contact; + const float t = + std::min(std::max((reduced - f_inner) * inverse_f_width, 0.0f), 1.0f); + const float t2 = t * t; + const float bridge = + t2 * t2 * (35.0f + t * (-84.0f + t * (70.0f - 20.0f * t))); + const float bump = t * (1.0f - t); + // dh/dr = h'(t) * dt/du * du/dr, the last factor being 1 / s. + const float bridge_slope = + 140.0f * bump * bump * bump * inverse_f_width * inverse_contact; + const float rise = + t2 * t2 * t * (7.0f + t * (-14.0f + t * (10.0f - 2.5f * t))); + state.seen = + contact * (0.5f * (f_inner + f_outer) + (f_outer - f_inner) * rise); + state.chain = bridge; + float cutoff_slope = 0.0f; + const float cutoff = cutoff_factor(state.seen, rcut, &cutoff_slope); + state.chi = cutoff * bridge; + state.chi_slope = cutoff_slope * bridge * bridge + cutoff * bridge_slope; + return state; +} + // === Cartesian harmonics === /// Evaluate the real Cartesian harmonics of degrees one through `LMAX`. @@ -100,7 +187,7 @@ inline float envelope(float radius, float rcut, float* derivative) { /// enters an angular moment. Output component `l * l + m - 1` holds /// \f$B^{(l)}_m\f$. template -inline void harmonics(float x, float y, float z, float* basis) { +DPA4C_CPU_INLINE void harmonics(float x, float y, float z, float* basis) { basis[0] = x; basis[1] = y; basis[2] = z; @@ -144,7 +231,7 @@ inline void harmonics(float x, float y, float z, float* basis) { /// \param cotangent Harmonic cotangents in the layout of :func:`harmonics`. /// \param direction Receives \f$\sum_m \bar B_m \partial B_m/\partial u\f$. template -inline void harmonics_backward( +DPA4C_CPU_INLINE void harmonics_backward( float x, float y, float z, const float* cotangent, float* direction) { float dx = cotangent[0]; float dy = cotangent[1]; @@ -213,10 +300,10 @@ inline void harmonics_backward( /// /// The interval stores six coefficient vectors per block, so the evaluation /// is one Horner chain of contiguous fused multiply-adds. -inline void spline_value(const float* __restrict interval, - float dx, - int blocks, - float* __restrict value) { +DPA4C_CPU_INLINE void spline_value(const float* __restrict interval, + float dx, + int blocks, + float* __restrict value) { for (int index = 0; index < blocks; ++index) { const float* __restrict coefficients = interval + index * 6 * kBlock; float* __restrict out = value + index * kBlock; @@ -236,11 +323,12 @@ inline void spline_value(const float* __restrict interval, /// Value and slope are two independent Horner chains over the same six /// coefficient vectors, so the loads are shared and the two dependency /// chains interleave. -inline void spline_value_and_derivative(const float* __restrict interval, - float dx, - int blocks, - float* __restrict value, - float* __restrict derivative) { +DPA4C_CPU_INLINE void spline_value_and_derivative( + const float* __restrict interval, + float dx, + int blocks, + float* __restrict value, + float* __restrict derivative) { for (int index = 0; index < blocks; ++index) { const float* __restrict coefficients = interval + index * 6 * kBlock; float* __restrict out_value = value + index * kBlock; @@ -268,10 +356,10 @@ inline void spline_value_and_derivative(const float* __restrict interval, } /// Evaluate the `R` shared mode profiles of one interval. -inline void mode_value(const float* __restrict modes, - float dx, - int count, - float* __restrict value) { +DPA4C_CPU_INLINE void mode_value(const float* __restrict modes, + float dx, + int count, + float* __restrict value) { for (int mode = 0; mode < count; ++mode) { const float* __restrict coefficients = modes + mode * 6; float accumulator = coefficients[5]; @@ -283,11 +371,11 @@ inline void mode_value(const float* __restrict modes, } /// Evaluate the mode profiles and their distance derivatives. -inline void mode_value_and_derivative(const float* __restrict modes, - float dx, - int count, - float* __restrict value, - float* __restrict derivative) { +DPA4C_CPU_INLINE void mode_value_and_derivative(const float* __restrict modes, + float dx, + int count, + float* __restrict value, + float* __restrict derivative) { for (int mode = 0; mode < count; ++mode) { const float* __restrict coefficients = modes + mode * 6; float accumulator = coefficients[5]; @@ -306,20 +394,61 @@ inline void mode_value_and_derivative(const float* __restrict modes, /// Everything one edge contributes, resolved once per direction. struct EdgeGeometry { float direction[3]; + /// True regularized length in Å, which the direction, the analytical pair + /// potential and the bridging window read. float radius; float chi; float chi_slope; + /// Slope of the length the table reads with respect to the true length. + float chain; int64_t interval; float dx; int64_t pair; }; +/// Evaluate the screened Coulomb series of one edge and its radial slope. +/// +/// Energy and slope share their exponentials: +/// \f$V=r^{-1}\sum_k A_k e^{-c_k r}\f$ and +/// \f$V'=-r^{-1}\sum_k A_k(c_k+r^{-1})e^{-c_k r}\f$. The exponent is held +/// at eighty, where a term lies at least nineteen orders of magnitude below +/// the float32 resolution of the sum it enters for every element pair within +/// 6 Å, and fourteen within 10 Å; beyond it the exponential would reach the +/// subnormal range, whose arithmetic stalls the core. +/// +/// \param row The constants `[A_0..A_3, c_0..c_3]` of the ordered type pair. +/// \param radius Regularized distance in Å. +/// \param slope Receives dV/dr in eV/Å. +/// \return The pair energy V in eV. +DPA4C_CPU_INLINE float pair_series(const float* row, + float radius, + float* slope) { + const float inverse = 1.0f / radius; + float energy = 0.0f; + float derivative = 0.0f; + for (int term = 0; term < 4; ++term) { + const float rate = row[4 + term]; + const float decayed = + row[term] * std::exp(-std::min(rate * radius, 80.0f)) * inverse; + energy += decayed; + derivative -= decayed * (rate + inverse); + } + *slope = derivative; + return energy; +} + /// Resolve one edge, returning false when it contributes nothing. -inline bool resolve_edge(const Arguments& arguments, - const Layout& layout, - int64_t edge, - int64_t center_type, - EdgeGeometry* geometry) { +/// +/// A padding neighbor never contributes. In a bridged specialization an edge +/// whose envelope vanishes, beyond the cutoff or inside the bridging window, +/// still resolves: it leaves the descriptor, which the scans read off +/// `chi == 0`, but an analytical pair potential keeps acting on it. +template +DPA4C_CPU_INLINE bool resolve_edge(const Arguments& arguments, + const Layout& layout, + int64_t edge, + int64_t center_type, + EdgeGeometry* geometry) { const int64_t neighbor = arguments.source[edge]; const int64_t neighbor_type = arguments.atype[neighbor]; if (neighbor_type >= layout.type_count - 1) { @@ -329,9 +458,13 @@ inline bool resolve_edge(const Arguments& arguments, const float squared = vector[0] * vector[0] + vector[1] * vector[1] + vector[2] * vector[2]; const float radius = std::sqrt(squared + arguments.eps * arguments.eps); - float slope = 0.0f; - const float chi = envelope(radius, arguments.rcut, &slope); - if (chi == 0.0f) { + const float contact = BRIDGED ? arguments.contact_radius[center_type] + + arguments.contact_radius[neighbor_type] + : 0.0f; + const EdgeEnvelope state = + envelope(radius, contact, arguments.rcut, arguments.f_inner, + arguments.f_outer, arguments.inverse_f_width); + if (!BRIDGED && state.chi == 0.0f) { return false; } const float inverse = 1.0f / radius; @@ -339,9 +472,10 @@ inline bool resolve_edge(const Arguments& arguments, geometry->direction[1] = vector[1] * inverse; geometry->direction[2] = vector[2] * inverse; geometry->radius = radius; - geometry->chi = chi; - geometry->chi_slope = slope; - const float coordinate = std::min(radius, arguments.table_max); + geometry->chi = state.chi; + geometry->chi_slope = state.chi_slope; + geometry->chain = state.chain; + const float coordinate = std::min(state.seen, arguments.table_max); int64_t interval = static_cast(coordinate / arguments.table_stride); interval = std::min(interval, layout.spline_count - 1); geometry->interval = interval; @@ -356,23 +490,40 @@ inline bool resolve_edge(const Arguments& arguments, #include "graph_compress_cpu_readout.inc" #include "graph_compress_cpu_scan.inc" -/// Return the entry points of this instruction-set level. -Kernels kernels(int lmax, bool has_modes) { +namespace { + +template +Kernels degree_kernels(bool has_modes) { + return has_modes ? Kernels{forward_scan, + backward_scan} + : Kernels{forward_scan, + backward_scan}; +} + +template +Kernels bridging_kernels(int lmax, bool has_modes) { switch (lmax) { case 2: - return has_modes - ? Kernels{forward_scan<2, true>, backward_scan<2, true>} - : Kernels{forward_scan<2, false>, backward_scan<2, false>}; + return degree_kernels<2, BRIDGED>(has_modes); case 3: - return has_modes - ? Kernels{forward_scan<3, true>, backward_scan<3, true>} - : Kernels{forward_scan<3, false>, backward_scan<3, false>}; + return degree_kernels<3, BRIDGED>(has_modes); default: - return has_modes - ? Kernels{forward_scan<4, true>, backward_scan<4, true>} - : Kernels{forward_scan<4, false>, backward_scan<4, false>}; + return degree_kernels<4, BRIDGED>(has_modes); } } +} // namespace + +/// Return the entry points of this instruction-set level. +/// +/// Zone bridging is a compile-time specialization like the angular degree and +/// the mode residual: a plain descriptor must not carry the switch arithmetic +/// of the edge envelope, nor the accumulator and the branch of the analytical +/// pair potential, of a bridged one. +Kernels kernels(int lmax, bool has_modes, bool bridged) { + return bridged ? bridging_kernels(lmax, has_modes) + : bridging_kernels(lmax, has_modes); +} + } // namespace DPA4C_CPU_ISA } // namespace deepmd_dpa4c_cpu diff --git a/source/op/pt/dpa4c/graph_compress_cpu_scan.inc b/source/op/pt/dpa4c/graph_compress_cpu_scan.inc index 634feef6a8..fb08556db5 100644 --- a/source/op/pt/dpa4c/graph_compress_cpu_scan.inc +++ b/source/op/pt/dpa4c/graph_compress_cpu_scan.inc @@ -104,7 +104,8 @@ class ScanWorkspace { }; /// Resolve the payload slot of one CSR entry. -inline int64_t payload_slot(const Arguments& arguments, int64_t entry) { +DPA4C_CPU_INLINE int64_t payload_slot(const Arguments& arguments, + int64_t entry) { return arguments.destination_order == nullptr ? entry : arguments.destination_order[entry]; } @@ -119,12 +120,12 @@ inline int64_t payload_slot(const Arguments& arguments, int64_t entry) { /// mode-carrying grades. The amplitude fits L1 several times over, so the /// repeated passes cost nothing the reload does not already hide. template -inline void film_amplitude(const Arguments& arguments, - const Layout& layout, - int64_t pair, - const float* __restrict value, - const float* __restrict modes, - float* __restrict amplitude) { +DPA4C_CPU_INLINE void film_amplitude(const Arguments& arguments, + const Layout& layout, + int64_t pair, + const float* __restrict value, + const float* __restrict modes, + float* __restrict amplitude) { const int padded = layout.padded_channels; const float* __restrict scale = arguments.tables->film.data() + pair * 2 * padded; const float* __restrict shift = scale + padded; @@ -150,7 +151,7 @@ inline void film_amplitude(const Arguments& arguments, } /// Accumulate every degree-wise moment of one node range. -template +template void forward_scan(const Arguments& arguments, const Layout& layout, int64_t node_begin, @@ -177,7 +178,9 @@ void forward_scan(const Arguments& arguments, continue; } EdgeGeometry geometry; - if (!resolve_edge(arguments, layout, edge, center_type, &geometry)) { + if (!resolve_edge(arguments, layout, edge, center_type, + &geometry) || + (BRIDGED && geometry.chi == 0.0f)) { continue; } const float* interval = @@ -248,7 +251,7 @@ void forward_scan(const Arguments& arguments, } /// Emit the edge cotangent of one node range. -template +template void backward_scan(const Arguments& arguments, const Layout& layout, int64_t node_begin, @@ -297,19 +300,37 @@ void backward_scan(const Arguments& arguments, const int64_t center_type = arguments.atype[node]; const bool real_center = center_type < layout.type_count - 1; + // Only a bridged specialization carries the analytical pair potential, + // and its descriptor-only backward still runs without one. Half of every + // pair energy belongs to the destination, so the half and the node's + // energy cotangent fold into one weight of the radial slope. + const bool has_pair = BRIDGED && arguments.pair_table != nullptr; + const float pair_weight = + has_pair ? 0.5f * static_cast(arguments.pair_seed[node]) : 0.0f; + float pair_energy = 0.0f; for (int64_t entry = arguments.row_ptr[node]; entry < arguments.row_ptr[node + 1]; ++entry) { const int64_t edge = payload_slot(arguments, entry); float* target = arguments.edge_gradient + 3 * edge; EdgeGeometry geometry; - const bool active = + const bool present = real_center && (arguments.edge_mask == nullptr || arguments.edge_mask[edge]) && - resolve_edge(arguments, layout, edge, center_type, &geometry); - if (!active) { - target[0] = 0.0f; - target[1] = 0.0f; - target[2] = 0.0f; + resolve_edge(arguments, layout, edge, center_type, + &geometry); + float pair_radial = 0.0f; + if (present && has_pair) { + float slope = 0.0f; + const float* row = arguments.pair_table + 8 * geometry.pair; + pair_energy += 0.5f * pair_series(row, geometry.radius, &slope); + pair_radial = pair_weight * slope; + } + if (!present || (BRIDGED && geometry.chi == 0.0f)) { + // Only the pair potential acts on an edge outside the descriptor. + for (int component = 0; component < 3; ++component) { + target[component] = + present ? pair_radial * geometry.direction[component] : 0.0f; + } continue; } const float* interval = @@ -444,11 +465,14 @@ void backward_scan(const Arguments& arguments, radial_lanes[lane] += value * film[channel] * slope[channel]; } } + // The table terms differentiate the length the table reads and are + // folded once through its slope; the envelope terms and the pair + // potential act on the true length. float envelope_accumulator = 0.0f; - float radial = 0.0f; + float table_gradient = 0.0f; for (int lane = 0; lane < kBlock; ++lane) { envelope_accumulator += envelope_lanes[lane]; - radial += radial_lanes[lane]; + table_gradient += radial_lanes[lane]; } if (HAS_MODES) { const float* __restrict mixing = @@ -470,13 +494,15 @@ void backward_scan(const Arguments& arguments, for (int lane = 0; lane < kBlock; ++lane) { mode_projection += mode_lanes[lane]; } - radial += mode_projection * mode_slope[mode]; + table_gradient += mode_projection * mode_slope[mode]; } } - radial += (envelope_accumulator + 2.0f * chi * divisor_cotangent[0] + - 4.0f * chi * chi_squared * divisor_cotangent[1] + - 2.0f * chi * amplitude_dot) * - geometry.chi_slope; + const float radial = + pair_radial + table_gradient * geometry.chain + + (envelope_accumulator + 2.0f * chi * divisor_cotangent[0] + + 4.0f * chi * chi_squared * divisor_cotangent[1] + + 2.0f * chi * amplitude_dot) * + geometry.chi_slope; float direction_cotangent[3]; harmonics_backward(geometry.direction[0], geometry.direction[1], @@ -493,6 +519,9 @@ void backward_scan(const Arguments& arguments, radial_part * geometry.direction[component]); } } + if (has_pair) { + arguments.pair_energy[node] += static_cast(pair_energy); + } } } diff --git a/source/op/pt/dpa4c/graph_compress_kernel.cuh b/source/op/pt/dpa4c/graph_compress_kernel.cuh index fa486baa14..a7184af891 100644 --- a/source/op/pt/dpa4c/graph_compress_kernel.cuh +++ b/source/op/pt/dpa4c/graph_compress_kernel.cuh @@ -79,12 +79,15 @@ template __global__ __launch_bounds__( Profile::Threads, - Profile:: - MinBlocksPerSm) void forward_kernel(Arguments args) { + edge_resident_blocks()) void forward_kernel(Arguments args) { using P = Profile; constexpr int EdgeWidth = P::ForwardEdgeWidth; constexpr int Groups = kWarpSize / EdgeWidth; @@ -111,6 +114,10 @@ __global__ __launch_bounds__( const auto* destination_order = static_cast(args.destination_order); const int center_type = static_cast(args.atype[args.node_begin + node]); + // The centre half of every pair length scale is a block constant, like the + // centre type it follows. + const float center_radius = + Bridged ? __ldg(args.contact_radius + center_type) : 0.0f; const long begin = args.destination_row_ptr[node]; const long end = args.destination_row_ptr[node + 1]; const int radial_modes = HasModes ? args.radial_modes : 0; @@ -144,17 +151,20 @@ __global__ __launch_bounds__( // broadcasting it from a leader. The addresses are identical inside the // group, so the memory system serves one transaction either way, whereas // a leader-only branch pays the same issue slots and adds ten shuffles. - const EdgeGeometry geometry = load_geometry( - edge, args.rcut, args.eps, args.edge_vec, edge_index, args.atype); - const TableLocation location = - locate_table(geometry.radius, args.table_stride, - args.table_max, args.interval_count); + const long source = static_cast(edge_index[edge]); + const int source_type = static_cast(args.atype[source]); if constexpr (!Canonical) { if (center_type >= args.type_count - 1 || - geometry.source_type >= args.type_count - 1) { + source_type >= args.type_count - 1) { continue; } } + const EdgeGeometry geometry = load_geometry( + edge, source, source_type, args.rcut, args.f_inner, args.f_outer, + args.inverse_f_width, center_radius, args.contact_radius, args.eps, + args.edge_vec); + const TableLocation location = locate_table( + geometry.seen, args.table_stride, args.table_max, args.interval_count); const TableRow row = table_row(args.table, location, args.table_width); const float coordinate = location.coordinate; if constexpr (HasModes) { @@ -1459,12 +1469,15 @@ template __global__ __launch_bounds__( Profile::Threads, - Profile:: - MinBlocksPerSm) void edge_backward_kernel(Arguments args) { + edge_resident_blocks()) void edge_backward_kernel(Arguments args) { using P = Profile; constexpr int EdgeWidth = P::BackwardEdgeWidth; constexpr int Groups = kWarpSize / EdgeWidth; @@ -1486,6 +1499,10 @@ __global__ __launch_bounds__( const auto* destination_order = static_cast(args.destination_order); const int center_type = static_cast(args.atype[args.node_begin + node]); + // The centre half of every pair length scale is a block constant, like the + // centre type it follows. + const float center_radius = + Bridged ? __ldg(args.contact_radius + center_type) : 0.0f; const long begin = args.destination_row_ptr[node]; const long end = args.destination_row_ptr[node + 1]; const int radial_modes = HasModes ? args.radial_modes : 0; @@ -1494,6 +1511,14 @@ __global__ __launch_bounds__( __ldg(args.moment_gradient + gradient_offset + P::MomentWidth); const float angular_mass_gradient = __ldg(args.moment_gradient + gradient_offset + P::MomentWidth + 1); + // Only a bridged specialization carries the analytical pair potential, and + // its descriptor-only backward still runs without one. Half of every pair + // energy belongs to the destination, so the half and the node's energy + // cotangent fold into one weight of the radial slope. + const bool has_pair = Bridged && args.pair_table != nullptr; + const float pair_weight = + has_pair ? 0.5f * static_cast(args.pair_seed[node]) : 0.0f; + float pair_energy = 0.0f; // The scalar cotangent is one vector of width C0 that every edge rereads. // Holding it in registers costs one entry per channel tile and spills the @@ -1586,14 +1611,11 @@ __global__ __launch_bounds__( } continue; } - const EdgeGeometry geometry = load_geometry( - edge, args.rcut, args.eps, args.edge_vec, edge_index, args.atype); - const TableLocation location = - locate_table(geometry.radius, args.table_stride, - args.table_max, args.interval_count); + const long source = static_cast(edge_index[edge]); + const int source_type = static_cast(args.atype[source]); if constexpr (!Canonical) { if (center_type >= args.type_count - 1 || - geometry.source_type >= args.type_count - 1) { + source_type >= args.type_count - 1) { if (thread == leader) { args.edge_gradient[edge * 3 + 0] = 0.0f; args.edge_gradient[edge * 3 + 1] = 0.0f; @@ -1607,6 +1629,25 @@ __global__ __launch_bounds__( continue; } } + const EdgeGeometry geometry = load_geometry( + edge, source, source_type, args.rcut, args.f_inner, args.f_outer, + args.inverse_f_width, center_radius, args.contact_radius, args.eps, + args.edge_vec); + const TableLocation location = locate_table( + geometry.seen, args.table_stride, args.table_max, args.interval_count); + // The pair term depends on the geometry alone, so it starts here and runs + // alongside the channel scan; its slope joins the radial cotangent at the + // tail, which keeps the exponentials off the dependency chain of the scan. + float pair_slope = 0.0f; + if (has_pair) { + const long pair_row = static_cast(center_type) * args.type_count + + geometry.source_type; + const PairTerm term = + pair_term(args.pair_table + pair_row * 8, base_channel, + geometry.radius, geometry.inverse_radius); + pair_energy = fmaf(0.5f, term.energy, pair_energy); + pair_slope = term.slope; + } const TableRow row = table_row(args.table, location, args.table_width); const float coordinate = location.coordinate; const bool clamped = location.clamped; @@ -1869,9 +1910,13 @@ __global__ __launch_bounds__( fmaf(4.0f * squared * envelope, angular_mass_gradient, fmaf(2.0f * envelope, scalar_mass_gradient, envelope_gradient)); } - radial_gradient = fmaf(envelope_gradient, - c3_envelope_derivative(geometry.radius, args.rcut), - radial_gradient); + // The table terms above differentiate the length the table reads; the + // pair potential acts on the true length and joins after the chain rule. + radial_gradient = radial_chain(geometry, args, center_radius, + radial_gradient, envelope_gradient); + if constexpr (Bridged) { + radial_gradient = fmaf(pair_weight, pair_slope, radial_gradient); + } // The basis VJP is linear in the radial and angular cotangents, so // applying it per lane reduces three Cartesian components instead of the @@ -1910,6 +1955,14 @@ __global__ __launch_bounds__( } } } + if (has_pair) { + // The lanes hold disjoint terms of disjoint edge groups, so one reduction + // over the warp is the pair energy of the node. + const float total = warp_sum(pair_energy); + if (thread == 0) { + args.pair_energy[node] += static_cast(total); + } + } } template @@ -1941,12 +1994,14 @@ template struct ForwardLauncher { static void run(const Arguments& args, cudaStream_t stream) { using P = Profile; - forward_kernel + forward_kernel <<(args.node_count), P::Threads, 0, stream>>>(args); } }; @@ -1955,6 +2010,7 @@ template struct BackwardLauncher { @@ -2015,7 +2071,8 @@ struct BackwardLauncher { } else { launch_node_backward(args, stream); } - edge_backward_kernel + edge_backward_kernel <<(args.node_count), P::Threads, 0, stream>>>(args); // The reserved edge slots beyond the physical count are only known on the // device, so the grid is sized from the storage bound and the surplus @@ -2043,26 +2100,47 @@ template class L> + bool Bridged, + template class L> void dispatch_topology(const Arguments& args, cudaStream_t stream) { const bool wide = args.index_kind == IndexKind::Bits64; if (args.canonical) { if (wide) { - L::run(args, stream); - } else { - L::run(args, + L::run(args, stream); + } else { + L::run( + args, stream); } } else { if (wide) { - L::run(args, stream); - } else { - L::run(args, + L::run(args, stream); + } else { + L::run( + args, stream); } } } +// Zone bridging is a compile-time specialization for the same reason as the +// native spin: a plain descriptor must not carry the switch arithmetic of the +// edge envelope, nor the accumulator and the branch of the analytical pair +// potential, of a bridged one. +template class L> +void dispatch_bridging(const Arguments& args, cudaStream_t stream) { + if (args.bridged) { + dispatch_topology(args, stream); + } else { + dispatch_topology(args, + stream); + } +} + // Native spin is a compile-time specialization for the same reason as the // mode residual: a spinless descriptor must not carry the spin accumulators, // the staged spin cotangent, or the wider moment state of one that has them. @@ -2071,12 +2149,12 @@ void dispatch_topology(const Arguments& args, cudaStream_t stream) { template class L> + template class L> void dispatch_spin(const Arguments& args, cudaStream_t stream) { if (args.has_spin) { - dispatch_topology(args, stream); + dispatch_bridging(args, stream); } else { - dispatch_topology(args, stream); + dispatch_bridging(args, stream); } } @@ -2085,7 +2163,7 @@ void dispatch_spin(const Arguments& args, cudaStream_t stream) { // vector temporaries and the shared profile cache of one that has them. template class L> + template class L> void dispatch_modes(const Arguments& args, cudaStream_t stream) { if (args.radial_modes > 0) { dispatch_spin(args, stream); @@ -2102,7 +2180,8 @@ void dispatch_modes(const Arguments& args, cudaStream_t stream) { // dispatch. The unreachable default is still checked rather than folded into // the highest degree, so that a degree outside the compiled set can only ever // fail loudly instead of running a kernel for a different model. -template class L> +template class L> void dispatch_degree(const Arguments& args, cudaStream_t stream) { switch (args.lmax) { case 2: diff --git a/source/op/pt/dpa4c/graph_compress_launch.h b/source/op/pt/dpa4c/graph_compress_launch.h index 70e52f1bd6..8c220f37d6 100644 --- a/source/op/pt/dpa4c/graph_compress_launch.h +++ b/source/op/pt/dpa4c/graph_compress_launch.h @@ -49,6 +49,17 @@ struct Arguments { float rcut = 0.0f; float eps = 0.0f; float degree_floor = 0.0f; + // Inner bridging switch of the edge envelope, as the two dimensionless + // fractions of a pair's own length scale that bound its window and the + // reciprocal of their difference. A bridged launch without a window keeps + // the outer fraction at zero, which no edge lies below. + float f_inner = 0.0f; + float f_outer = 0.0f; + float inverse_f_width = 0.0f; + // Whether the launch belongs to a zone-bridging composition, that is whether + // it carries a bridging window or an analytical pair potential. Only the + // bridged specializations evaluate either. + bool bridged = false; bool canonical = false; // Whether the native spin branch is active. The neighbour spin width is // derived from the degree profile, so presence is the whole runtime choice. @@ -71,6 +82,11 @@ struct Arguments { const float* coupling_value = nullptr; const float* output_mean = nullptr; const float* output_inv_std = nullptr; + // Per-element length scale in Å, indexed by atom type over the whole type + // table including the padding row. The sum of the two entries of a pair is + // the distance the window fractions are measured against. Read by the + // bridged specializations only. + const float* contact_radius = nullptr; const float* descriptor_gradient = nullptr; const float* state = nullptr; @@ -85,6 +101,18 @@ struct Arguments { const float* spin_pair = nullptr; const float* spin_type = nullptr; + // === Analytical pair potential === + // Present together or not at all, and read by the edge backward only, which + // is the one scan that sees an edge together with the energy cotangent of + // its destination. ``pair_table`` holds, per ordered type pair, the four + // amplitudes and the four decay rates of the screened Coulomb series + // V(r) = (1/r) sum_k A_k exp(-c_k r). ``pair_seed`` is that cotangent and + // ``pair_energy`` the fp64 node energies onto which the scan adds V/2 per + // edge. Both are node indexed and offset by the caller, like ``descriptor``. + const float* pair_table = nullptr; + const double* pair_seed = nullptr; + double* pair_energy = nullptr; + float* descriptor = nullptr; float* state_out = nullptr; float* moment_gradient = nullptr; @@ -364,23 +392,46 @@ struct Profile { NodeLanes != 0 ? NodeLanes : kNodeLanesNarrow; static constexpr int NodeGroups = kWarpSize / NodeWidth; static constexpr int Threads = kWarpSize; - - // Resident blocks the edge kernels are compiled for. A block is one warp, so - // thirty-two of them exhaust the 65,536-register file at sixty-four - // registers per thread, and a lower target raises the per-thread budget in - // proportion at the cost of occupancy. - // - // Degree two holds its geometric working set in that budget and spills only - // what the five spin families add, so twenty-four blocks buy eighty - // registers, remove the spill outright and are worth 2.0% to 4.5% of the - // step. Degree three overflows the budget on geometry alone -- it spills - // without the spin branch and still spills at eighty and at ninety-six - // registers -- so the occupancy it would give up buys an incomplete fix and - // costs 7.4% to 9.6%. The relief is therefore taken only where it is - // complete. - static constexpr int MinBlocksPerSm = (HasSpin && Lmax == 2) ? 24 : 32; }; +/// Resident blocks an edge kernel is compiled for, passed as the minimum block +/// count of its launch bounds. +/// +/// A block is one warp, so the count sets the register budget per thread: +/// thirty-two blocks exhaust the 65,536-register file at sixty-four registers +/// per thread, and a lower count raises the budget in proportion at the cost of +/// occupancy. Zero sets no budget and leaves the register count to the +/// compiler. +template +constexpr int edge_resident_blocks() { + // Measured on compute capability 9.0, degree two holds its geometric working + // set in sixty-four registers and spills only what the five spin families + // add, so twenty-four blocks buy eighty registers, remove the spill outright + // and are worth 2.0% to 4.5% of the step. Degree three overflows the budget + // on geometry alone -- it spills without the spin branch and still spills at + // eighty and at ninety-six registers -- so the occupancy it would give up + // buys an incomplete fix and costs 7.4% to 9.6%. The relief is therefore + // taken only where it is complete. + if (HasSpin && Lmax == 2) { + return 24; + } +#if defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ == 890 || (__CUDA_ARCH__ >= 1200 && __CUDA_ARCH__ < 1300)) + // A multiprocessor of compute capability 8.9 or 12.x holds at most + // twenty-four blocks, so eighty registers is the tightest budget there. It + // serves the radial-mode kernels of sixty-four channels and more, to which + // the compiler gives 86 to 128 registers: the budget restores the resident + // warps those counts cost for a spill of at most 80 bytes and shortens the + // plain kernels by 3% to 13% on 12.0. The other kernels run fastest on the + // compiler's own count, which keeps the spin families in registers and suits + // the narrow profiles, whose two-lane edge groups keep sixteen edges in + // flight per warp. + return HasModes && Channels >= 64 ? 24 : 0; +#else + return 32; +#endif +} + /// Scalar widths that own a compiled specialization. #define DPA4C_FOR_EACH_CHANNEL(macro) \ macro(8) macro(16) macro(32) macro(64) macro(128) diff --git a/source/op/pt/dpa4c/ops.cc b/source/op/pt/dpa4c/ops.cc index 010f15fd49..4569bb9bbe 100644 --- a/source/op/pt/dpa4c/ops.cc +++ b/source/op/pt/dpa4c/ops.cc @@ -15,6 +15,20 @@ // permutation alongside the row pointers. The canonical one is the compact // deployment ABI: destination-major payload, identity permutation, no mask, // and source indices only. +// +// The two operators that close an energy-force evaluation, the generic +// backward and the fused canonical energy-gradient, also evaluate an +// analytical pair potential when `pair_table` is not empty. Its energy and +// radial slope share their exponentials and need no saved state, so the edge +// scan that emits the edge gradient produces both: it adds +// `pair_seed_i * dV/dr / 2` to the radial cotangent of every edge of node `i` +// and accumulates `V / 2` into the node energy (`pair_energy`, or the +// `energy` of the fused operator, whose `seed` is the pair seed). +// +// Every schema carries the zone-bridging window as the two dimensionless +// fractions `f_inner` and `f_outer` of a pair's own length scale, together +// with the per-element radii `contact_radius` whose pairwise sums are those +// length scales in Å. Equal fractions denote a descriptor without a window. #include @@ -26,10 +40,10 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, " "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, " "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, " - "Tensor spin, Tensor spin_pair, Tensor spin_type, " + "Tensor spin, Tensor spin_pair, Tensor spin_type, Tensor contact_radius, " "bool canonical, int lmax, float table_stride, float table_max, " - "float rcut, float eps, float degree_floor) " - "-> (Tensor descriptor, Tensor state)"); + "float rcut, float eps, float degree_floor, float f_inner, " + "float f_outer) -> (Tensor descriptor, Tensor state)"); library.def( "dpa4c_graph_compress_backward(Tensor descriptor_gradient, " "Tensor state, Tensor edge_vec, Tensor edge_index, Tensor edge_mask, " @@ -38,19 +52,22 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "Tensor type_embedding, Tensor readout_matrices, Tensor coupling_meta, " "Tensor coupling_entry, Tensor coupling_value, Tensor output_mean, " "Tensor output_inv_std, Tensor spin, Tensor spin_pair, " - "Tensor spin_type, bool canonical, int lmax, float table_stride, " - "float table_max, float rcut, float eps, float degree_floor) " + "Tensor spin_type, Tensor contact_radius, bool canonical, int lmax, " + "float table_stride, float table_max, float rcut, float eps, " + "float degree_floor, float f_inner, float f_outer, Tensor pair_table, " + "Tensor pair_seed) " "-> (Tensor edge_gradient, Tensor spin_gradient, " - "Tensor edge_spin_gradient)"); + "Tensor edge_spin_gradient, Tensor pair_energy)"); library.def( "dpa4c_canonical_compress(Tensor edge_vec, Tensor source, " "Tensor destination_row_ptr, Tensor atype, Tensor table, " "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, " "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, " "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, " - "Tensor spin, Tensor spin_pair, Tensor spin_type, " + "Tensor spin, Tensor spin_pair, Tensor spin_type, Tensor contact_radius, " "int lmax, float table_stride, float table_max, float rcut, float eps, " - "float degree_floor) -> (Tensor descriptor, Tensor state)"); + "float degree_floor, float f_inner, float f_outer) " + "-> (Tensor descriptor, Tensor state)"); library.def( "dpa4c_canonical_compress_backward(Tensor descriptor_gradient, " "Tensor state, Tensor edge_vec, Tensor source, " @@ -58,9 +75,9 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, " "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, " "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, " - "Tensor spin, Tensor spin_pair, Tensor spin_type, " + "Tensor spin, Tensor spin_pair, Tensor spin_type, Tensor contact_radius, " "int lmax, float table_stride, float table_max, float rcut, float eps, " - "float degree_floor) " + "float degree_floor, float f_inner, float f_outer) " "-> (Tensor edge_gradient, Tensor spin_gradient, " "Tensor edge_spin_gradient)"); library.def( @@ -70,9 +87,9 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, " "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, " "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, " - "Tensor spin, Tensor spin_pair, Tensor spin_type, " + "Tensor spin, Tensor spin_pair, Tensor spin_type, Tensor contact_radius, " "int lmax, float table_stride, float table_max, float rcut, float eps, " - "float degree_floor) " + "float degree_floor, float f_inner, float f_outer) " "-> (Tensor edge_gradient, Tensor spin_gradient, " "Tensor edge_spin_gradient)"); library.def( @@ -81,11 +98,12 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, " "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, " "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, " - "Tensor spin, Tensor spin_pair, Tensor spin_type, " + "Tensor spin, Tensor spin_pair, Tensor spin_type, Tensor contact_radius, " "int lmax, float table_stride, float table_max, float rcut, float eps, " - "float degree_floor, Tensor[] ws, Tensor[] bs, int[] resnets, " + "float degree_floor, float f_inner, float f_outer, " + "Tensor[] ws, Tensor[] bs, int[] resnets, " "Tensor w_head, Tensor b_head, Tensor bias_atom_e, int act, " - "Tensor seed, int tile) " + "Tensor seed, int tile, Tensor pair_table) " "-> (Tensor energy, Tensor edge_gradient, Tensor spin_gradient, " "Tensor edge_spin_gradient)"); } diff --git a/source/tests/common/dpmodel/test_atomic_model_capabilities.py b/source/tests/common/dpmodel/test_atomic_model_capabilities.py index 70b504353b..77b8b7530f 100644 --- a/source/tests/common/dpmodel/test_atomic_model_capabilities.py +++ b/source/tests/common/dpmodel/test_atomic_model_capabilities.py @@ -12,6 +12,9 @@ from deepmd.dpmodel.atomic_model import ( DPAtomicModel, ) +from deepmd.dpmodel.atomic_model.base_atomic_model import ( + BaseAtomicModel, +) from deepmd.dpmodel.atomic_model.inner_potential import ( InnerPotentialAtomicModel, ) @@ -69,7 +72,13 @@ def _dpa4_descriptor(bridging: bool) -> DescrptDPA4: "random_gamma": False, } if bridging: - kwargs.update(inner_clamp_r_inner=0.8, inner_clamp_r_outer=1.2) + # The absolute scale gives every pair the unit length scale, so the two + # fractions are the window radii in Å. + kwargs.update( + inner_clamp_f_inner=0.8, + inner_clamp_f_outer=1.2, + inner_clamp_scale="absolute", + ) return DescrptDPA4(**kwargs) @@ -91,7 +100,19 @@ def test_base_defaults(cap, default) -> None: def test_base_graph_edge_dtype_default() -> None: """float64 is the model-agnostic edge-geometry ABI.""" - assert _inner_potential().graph_edge_dtype() == "float64" + plain = _dp_atomic_model(_dpa4_descriptor(bridging=False)) + assert BaseAtomicModel.graph_edge_dtype(plain) == "float64" + + +def test_analytical_term_accepts_single_precision_edges() -> None: + """The learned sibling decides the dtype of a bridged composition. + + The analytical term evaluates on edges of either precision, so under the + rule of the composition (float32 only if every child accepts it) it + answers float32 and never vetoes a compressed single-precision sibling. + """ + assert _inner_potential().graph_edge_dtype() == "float32" + assert _inner_potential().fused_decomposition()[0] is None def test_dp_atomic_model_delegates_to_descriptor() -> None: @@ -118,6 +139,18 @@ def test_dp_atomic_model_delegates_to_descriptor() -> None: ) +def test_hybrid_descriptor_refuses_a_bridged_child() -> None: + """A concatenated descriptor has no fitting output for a child's readout gate.""" + from deepmd.dpmodel.descriptor.hybrid import ( + DescrptHybrid, + ) + + with pytest.raises(NotImplementedError, match="bridged child"): + DescrptHybrid( + [_dpa4_descriptor(bridging=True), _dpa4_descriptor(bridging=False)] + ) + + def test_hybrid_descriptor_aggregates_dense_comm() -> None: """A hybrid is only as comm-capable as its least capable child. diff --git a/source/tests/common/dpmodel/test_descriptor_dpa4c.py b/source/tests/common/dpmodel/test_descriptor_dpa4c.py index abe2b63b97..f75e3d859f 100644 --- a/source/tests/common/dpmodel/test_descriptor_dpa4c.py +++ b/source/tests/common/dpmodel/test_descriptor_dpa4c.py @@ -655,6 +655,166 @@ def test_allowed_degree_triples_and_gaunt_couplings(self) -> None: np.testing.assert_allclose(np.linalg.norm(coupling), 1.0, atol=1e-15) +class TestDPA4CBridgingWindow: + """The bridging window removes a close pair from the descriptor.""" + + R_INNER = 0.5 + R_OUTER = 0.8 + + def setup_method(self) -> None: + self.config = { + "rcut": 3.0, + "ntypes": 3, + "channels": 8, + "lmax": 3, + "n_radial": 4, + "radial_modes": 2, + "precision": "float64", + "seed": 17, + } + self.windowed = DescrptDPA4C( + **self.config, + inner_clamp_f_inner=self.R_INNER, + inner_clamp_f_outer=self.R_OUTER, + inner_clamp_scale="absolute", + ) + self.plain = DescrptDPA4C(**self.config) + + @staticmethod + def cluster(distance: float) -> tuple[np.ndarray, np.ndarray]: + """Place a type-2 atom `distance` from the type-0 origin atom. + + The direction keeps it more than 1.3 Å from every other atom, so the + origin pair is the only one the window acts on. + """ + direction = np.array([0.8, 0.36, -0.48]) + coord = COORD.copy() + coord[0, 1] = distance * direction + atype = np.array([[0, 2, 1, 1, 1]], dtype=np.int64) + return coord, atype + + @pytest.mark.parametrize("distance", [0.05, 0.3, 0.5]) + def test_frozen_pair_equals_removed_edge(self, distance: float) -> None: + coord, atype = self.cluster(distance) + removed = DescrptDPA4C(**self.config, exclude_types=[(0, 2), (2, 0)]) + windowed = evaluate(self.windowed, coord, atype) + # The plain descriptor must see the pair, or the comparison is vacuous. + assert np.abs(evaluate(self.plain, coord, atype) - windowed).max() > 1e-3 + np.testing.assert_array_equal(windowed, evaluate(removed, coord, atype)) + + def test_pairs_beyond_the_window_are_unchanged(self) -> None: + coord, atype = self.cluster(self.R_OUTER) + np.testing.assert_array_equal( + evaluate(self.windowed, coord, atype), + evaluate(self.plain, coord, atype), + ) + + def test_switch_rises_monotonically_across_the_window(self) -> None: + distances = np.linspace(self.R_INNER - 0.01, self.R_OUTER + 0.01, 65) + switch = [] + for distance in distances: + coord, atype = self.cluster(float(distance)) + graph, atype_local = build_graph(self.windowed, coord, atype) + edge = np.argmin(np.abs(np.linalg.norm(graph.edge_vec, axis=-1) - distance)) + switch.append( + edge_features(self.windowed, graph, atype_local)[2][edge] + / edge_features(self.plain, graph, atype_local)[2][edge] + ) + switch = np.asarray(switch) + inside = (distances > self.R_INNER) & (distances < self.R_OUTER) + assert np.all(switch[distances < self.R_INNER] == 0.0) + assert np.all(switch[distances > self.R_OUTER] == 1.0) + assert np.all(np.diff(switch[inside]) > 0.0) + # A C3 switch keeps its first three differences continuous across both + # radii, which bounds the fourth one by the fourth power of the spacing. + assert np.abs(np.diff(switch, n=4)).max() < 5.0e-4 + + def test_clamp_freezes_at_the_midpoint_and_opens_with_the_switch(self) -> None: + clamp = self.windowed.bridging_clamp + switch = self.windowed.bridging_switch + midpoint = 0.5 * (self.R_INNER + self.R_OUTER) + distances = np.linspace(0.05, self.R_OUTER + 0.2, 4001) + contact = np.ones_like(distances) + seen = clamp.call(distances, contact) + np.testing.assert_array_equal(seen[distances <= self.R_INNER], midpoint) + beyond = distances >= self.R_OUTER + np.testing.assert_array_equal(seen[beyond], distances[beyond]) + # The clamp is the integral of the switch, so its slope is the switch. + np.testing.assert_allclose( + np.gradient(seen, distances)[2:-2], + switch.call(distances, contact)[2:-2], + atol=1.0e-6, + ) + + @pytest.mark.parametrize("distance", [0.55, 0.65, 0.75]) + def test_edge_inside_the_window_reads_the_clamped_length( + self, distance: float + ) -> None: + """The radial features see the clamped length, the harmonics the true one.""" + contact = np.ones(1) + seen = float( + self.windowed.bridging_clamp.call(np.array([distance]), contact)[0] + ) + opening = float( + self.windowed.bridging_switch.call(np.array([distance]), contact)[0] + ) + assert seen > distance + features = {} + for name, descriptor, length in ( + ("windowed", self.windowed, distance), + ("plain", self.plain, seen), + ): + coord, atype = self.cluster(length) + graph, atype_local = build_graph(descriptor, coord, atype) + edge = np.argmin(np.abs(np.linalg.norm(graph.edge_vec, axis=-1) - length)) + features[name] = [ + block[edge] for block in edge_features(descriptor, graph, atype_local) + ] + amplitude, harmonics, envelope = features["windowed"] + plain_amplitude, plain_harmonics, plain_envelope = features["plain"] + np.testing.assert_allclose(amplitude, opening * plain_amplitude, rtol=1.0e-12) + np.testing.assert_allclose(envelope, opening * plain_envelope, rtol=1.0e-12) + np.testing.assert_allclose(harmonics, plain_harmonics, rtol=1.0e-12) + + def test_serialization_roundtrip_keeps_the_window(self) -> None: + coord, atype = self.cluster(0.65) + data = self.windowed.serialize() + assert data["@version"] == 2 + assert data["inner_clamp_f_inner"] == self.R_INNER + assert data["inner_clamp_f_outer"] == self.R_OUTER + assert data["inner_clamp_scale"] == "absolute" + np.testing.assert_array_equal( + evaluate(DescrptDPA4C.deserialize(data), coord, atype), + evaluate(self.windowed, coord, atype), + ) + + def test_a_record_without_a_window_restores_the_plain_descriptor(self) -> None: + """A record of the first format carries no window and reads as unbridged.""" + data = self.plain.serialize() + data["@version"] = 1 + for key in ("inner_clamp_f_inner", "inner_clamp_f_outer", "inner_clamp_scale"): + del data[key] + restored = DescrptDPA4C.deserialize(data) + assert restored.bridging_switch is None + coord, atype = self.cluster(0.3) + np.testing.assert_array_equal( + evaluate(restored, coord, atype), + evaluate(self.plain, coord, atype), + ) + + def test_window_is_branch_local_under_parameter_sharing(self) -> None: + self.windowed.share_params(self.plain, shared_level=0) + coord, atype = self.cluster(0.3) + assert self.windowed.radial_embedding is self.plain.radial_embedding + assert ( + np.abs( + evaluate(self.windowed, coord, atype) + - evaluate(self.plain, coord, atype) + ).max() + > 1e-3 + ) + + @pytest.mark.parametrize( ("config", "error"), [ @@ -672,6 +832,9 @@ def test_allowed_degree_triples_and_gaunt_couplings(self) -> None: ({"radial_modes": True}, ValueError), ({"radial_modes": -1}, ValueError), ({"spin": {}}, NotImplementedError), + ({"inner_clamp_f_inner": 0.5}, ValueError), + ({"inner_clamp_f_inner": 0.8, "inner_clamp_f_outer": 0.5}, ValueError), + ({"inner_clamp_scale": "bohr"}, ValueError), ], ) def test_configuration_boundaries( diff --git a/source/tests/common/dpmodel/test_descrpt_dpa4.py b/source/tests/common/dpmodel/test_descrpt_dpa4.py index f145091b34..015437a089 100644 --- a/source/tests/common/dpmodel/test_descrpt_dpa4.py +++ b/source/tests/common/dpmodel/test_descrpt_dpa4.py @@ -1,5 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Torch-free unit tests for the dpmodel DPA4 (SeZM) descriptor.""" +"""Unit tests for the dpmodel DPA4 (SeZM) descriptor. + +The suite runs on the dpmodel backend alone except for the source-gate +cross-backend check, whose pt import lives inside the test function because +ruff TID253 bans module-level ``deepmd.pt`` imports under ``source/tests``. +""" import numpy as np import pytest @@ -61,6 +66,25 @@ def make_descriptor(**overrides) -> DescrptDPA4: return DescrptDPA4(**kwargs) +# The absolute scale gives every element the unit radius, so each pair's length +# scale is 1 Å and the two fractions are the window radii in Å directly. +BRIDGING_WINDOW = { + "inner_clamp_f_inner": 0.5, + "inner_clamp_f_outer": 1.0, + "inner_clamp_scale": "absolute", +} + +# A hand-built graph of four nodes and six edges for the source gates, read +# against ``BRIDGING_WINDOW``: the absolute scale gives every pair the unit +# length scale, so an edge length is already its reduced distance. Node 0 emits +# one frozen edge (inside the inner radius, amplitude exactly zero) and one open +# one, node 1 two open edges, node 2 a single edge beyond the outer radius +# (amplitude exactly one) and node 3 a single open edge. +GATE_EDGE_LEN = np.array([[0.3], [0.7], [0.7], [0.9], [1.5], [0.6]], dtype=np.float64) +GATE_SRC = np.array([0, 0, 1, 1, 2, 3], dtype=np.int64) +GATE_N_NODES = 4 + + def make_inputs(seed=5, 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)) @@ -76,7 +100,8 @@ def test_shapes_and_interface(self) -> None: nf, nloc = atype.shape out = dd.call(coord.reshape(nf, -1), atype, nlist, mapping=None) assert out[0].shape == (nf, nloc, dd.get_dim_out()) - assert out[1:] == (None, None, None, None) + # rot_mat, g2, h2, sw, and the source gate of an unbridged descriptor + assert out[1:] == (None, None, None, None, None) assert np.isfinite(np.asarray(out[0])).all() # standard descriptor surface assert dd.get_rcut() == 4.0 @@ -106,7 +131,7 @@ def test_message_passing_semantics(self) -> None: 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) + dd_bridge = make_descriptor(**BRIDGING_WINDOW) assert dd_bridge.has_message_passing() is True assert dd_bridge.has_message_passing_across_ranks() is True @@ -116,7 +141,7 @@ def test_capability_split_needs_vs_supports(self) -> None: models too since the SFPG cross-rank completion -- issue #5906). """ dd_plain = make_descriptor() - dd_bridged = make_descriptor(inner_clamp_r_inner=0.5, inner_clamp_r_outer=1.0) + dd_bridged = make_descriptor(**BRIDGING_WINDOW) assert dd_plain.has_message_passing_across_ranks() is True assert dd_bridged.has_message_passing_across_ranks() is True assert dd_plain.supports_edge_parallel() is True @@ -126,40 +151,121 @@ def test_gate_partial_exchange_dpmodel_raises(self) -> None: """The dpmodel backend is the single-process reference; comm on a bridged model must raise, never silently compute a partial gate. """ - dd = make_descriptor(inner_clamp_r_inner=0.5, inner_clamp_r_outer=1.0) + dd = make_descriptor(**BRIDGING_WINDOW) with pytest.raises(NotImplementedError, match="dpmodel"): dd._gate_partial_exchange(np.zeros((4, 2)), {"nlocal": 2}) - def test_edge_src_gate_identity_exchange_is_noop(self) -> None: - """The hook seam: an identity exchange reproduces the no-hook gate - bit-exactly (pins the pack/unpack layout [log_eta, zero_count]). + def test_source_gates_are_the_full_and_leave_one_out_products(self) -> None: + """The node gate multiplies every pair of a node, the edge gate all but one. + + A pair inside the frozen zone contributes an amplitude of exactly + zero. It mutes its source node and every OTHER edge that node emits, + while the frozen edge itself keeps the product over the remaining + pairs of its source, so the two atoms of a frozen pair go on seeing + each other at the clamped distance. """ from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import ( - compute_edge_src_gate, + compute_source_gates, + ) + + switch = make_descriptor(**BRIDGING_WINDOW).bridging_switch + contact = np.ones_like(GATE_EDGE_LEN) + w = np.asarray(switch.call(GATE_EDGE_LEN, contact))[:, 0] + # the data spans all three regimes of the switch: frozen, transition, + # and fully open + assert w[0] == 0.0 + assert all(0.0 < w[i] < 1.0 for i in (1, 2, 3, 5)) + assert w[4] == 1.0 + + node_gate, edge_gate = compute_source_gates( + edge_len=GATE_EDGE_LEN, + edge_contact=contact, + src=GATE_SRC, + n_nodes=GATE_N_NODES, + bridging_switch=switch, ) + assert node_gate.shape == (GATE_N_NODES,) + assert edge_gate.shape == (GATE_EDGE_LEN.shape[0], 1) + + # node 0 owns the frozen pair, so its full product is exactly zero + assert node_gate[0] == 0.0 + np.testing.assert_allclose(node_gate[1], w[2] * w[3], rtol=1e-14, atol=0.0) + np.testing.assert_allclose(node_gate[2], w[4], rtol=1e-14, atol=0.0) + np.testing.assert_allclose(node_gate[3], w[5], rtol=1e-14, atol=0.0) + + gate = np.asarray(edge_gate)[:, 0] + # node 0: the frozen edge keeps its source's other pair, and that + # other edge is muted by the frozen one + np.testing.assert_allclose(gate[0], w[1], rtol=1e-14, atol=0.0) + assert gate[1] == 0.0 + # node 1: each of the two open edges keeps the other one + np.testing.assert_allclose(gate[2], w[3], rtol=1e-14, atol=0.0) + np.testing.assert_allclose(gate[3], w[2], rtol=1e-14, atol=0.0) + # nodes 2 and 3 emit one edge each: the leave-one-out product is empty + np.testing.assert_array_equal(gate[4:], np.ones(2)) + + def test_source_gates_match_the_pt_backend(self) -> None: + """Both backends compute the two gates through the same reduction. + + The log-sum decomposition is written the same way on either side, so + on the same float64 input the two implementations agree bit for bit. + """ + import torch - dd = make_descriptor(inner_clamp_r_inner=0.5, inner_clamp_r_outer=1.0) - sw = dd.bridging_switch - # 4 nodes, 5 edges; one edge inside r_inner (w=0, hard-freezes its - # src node), the rest spread across the transition zone and beyond. - el = np.array([[0.3], [0.7], [0.9], [1.5], [0.75]], dtype=np.float64) - src = np.array([0, 0, 1, 2, 3], dtype=np.int64) - gate_ref = compute_edge_src_gate( - edge_len=el, src=src, n_nodes=4, bridging_switch=sw + from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import ( + compute_source_gates, + ) + from deepmd.pt.model.descriptor.sezm_nn.edge_cache import ( + compute_source_gates as compute_source_gates_pt, ) - gate_hook = compute_edge_src_gate( - edge_len=el, - src=src, - n_nodes=4, - bridging_switch=sw, - node_partial_exchange=lambda p: p, + from deepmd.pt.model.descriptor.sezm_nn.radial import ( + BridgingSwitch as BridgingSwitchPT, ) - np.testing.assert_array_equal(gate_ref, gate_hook) - # geometry sanity: node 0 is hard-frozen, node 3 is inside the - # transition zone (0 < gate < 1) -- the test data actually - # exercises both gate branches. - assert gate_ref[0, 0] == 0.0 - assert 0.0 < gate_ref[4, 0] < 1.0 + + switch = make_descriptor(**BRIDGING_WINDOW).bridging_switch + contact = np.ones_like(GATE_EDGE_LEN) + node_gate, edge_gate = compute_source_gates( + edge_len=GATE_EDGE_LEN, + edge_contact=contact, + src=GATE_SRC, + n_nodes=GATE_N_NODES, + bridging_switch=switch, + ) + + switch_pt = BridgingSwitchPT(switch.f_inner, switch.f_outer).to("cpu") + node_gate_pt, edge_gate_pt = compute_source_gates_pt( + edge_len=torch.from_numpy(GATE_EDGE_LEN), + edge_contact=torch.from_numpy(contact), + src=torch.from_numpy(GATE_SRC), + n_nodes=GATE_N_NODES, + bridging_switch=switch_pt, + ) + np.testing.assert_array_equal(np.asarray(node_gate), node_gate_pt.numpy()) + np.testing.assert_array_equal(np.asarray(edge_gate), edge_gate_pt.numpy()) + + def test_source_gate_identity_exchange_is_noop(self) -> None: + """The hook seam: an identity exchange reproduces the no-hook gates + bit-exactly (pins the pack/unpack layout [log_eta, zero_count]). + """ + from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import ( + compute_source_gates, + ) + + switch = make_descriptor(**BRIDGING_WINDOW).bridging_switch + contact = np.ones_like(GATE_EDGE_LEN) + kwargs = { + "edge_len": GATE_EDGE_LEN, + "edge_contact": contact, + "src": GATE_SRC, + "n_nodes": GATE_N_NODES, + "bridging_switch": switch, + } + node_ref, edge_ref = compute_source_gates(**kwargs) + node_hook, edge_hook = compute_source_gates( + **kwargs, node_partial_exchange=lambda p: p + ) + np.testing.assert_array_equal(node_ref, node_hook) + np.testing.assert_array_equal(edge_ref, edge_hook) def test_serialize_roundtrip_exact(self) -> None: dd = make_descriptor() @@ -289,7 +395,7 @@ def test_legacy_spin_gate_is_squared_on_deserialize(self) -> None: """ dd = make_descriptor(use_spin=[True, False, False]) data = dd.serialize() - assert data["@version"] == DescrptDPA4.LATEST_VERSION == 1.2 + assert data["@version"] == DescrptDPA4.LATEST_VERSION == 1.3 data["@version"] = 1.1 data["@variables"]["env_seed_embedding.spin_scale"] = np.full( (1,), 3.0, dtype=np.float64 @@ -297,11 +403,11 @@ def test_legacy_spin_gate_is_squared_on_deserialize(self) -> None: migrated = DescrptDPA4.deserialize(data) np.testing.assert_allclose(migrated.env_seed_embedding.spin_scale, 9.0) - assert migrated.version == 1.2 + assert migrated.version == 1.3 reloaded = DescrptDPA4.deserialize(migrated.serialize()) np.testing.assert_allclose(reloaded.env_seed_embedding.spin_scale, 9.0) - assert reloaded.version == 1.2 + assert reloaded.version == 1.3 def test_legacy_spin_free_routes_are_zeroed_on_deserialize(self) -> None: data = make_descriptor(use_spin=[False, False, False]).serialize() @@ -327,7 +433,42 @@ def test_legacy_spin_free_routes_are_zeroed_on_deserialize(self) -> None: np.testing.assert_array_equal( variables[mag_layer1_key], np.full_like(variables[mag_layer1_key], 5.0) ) - assert migrated.version == 1.2 + assert migrated.version == 1.3 + + def test_bridged_records_before_the_window_are_refused(self) -> None: + """A bridged record below 1.3 was trained under other window mechanics. + + Its radii still translate (``migrate_inner_clamp_keys``), but no + rewrite of the stored variables expresses the clamp frozen at the + inner radius or the full-product gate it was trained with, so the + record is refused instead of being read under the current function. + An unbridged record of the same version loads. + """ + data = make_descriptor(**BRIDGING_WINDOW).serialize() + data["@version"] = 1.2 + config = data["config"] + config["inner_clamp_r_inner"] = config.pop("inner_clamp_f_inner") + config["inner_clamp_r_outer"] = config.pop("inner_clamp_f_outer") + config.pop("inner_clamp_scale") + with pytest.raises(ValueError, match="Retrain"): + DescrptDPA4.deserialize(data) + + plain = make_descriptor().serialize() + plain["@version"] = 1.2 + assert DescrptDPA4.deserialize(plain).version == 1.3 + + def test_legacy_plain_upgrade_preserves_the_descriptor(self) -> None: + """The 1.3 format transition is inert without a bridging window.""" + descriptor = make_descriptor() + data = descriptor.serialize() + data["@version"] = 1.2 + restored = DescrptDPA4.deserialize(data) + coord, atype, nlist = make_inputs() + nf = atype.shape[0] + expected = descriptor.call(coord.reshape(nf, -1), atype, nlist)[0] + actual = restored.call(coord.reshape(nf, -1), atype, nlist)[0] + np.testing.assert_array_equal(actual, expected) + assert restored.serialize()["@version"] == 1.3 def test_pre_spin_versions_keep_their_own_tag(self) -> None: """Version 1.0 predates the spin route and the 1.1 forward math. diff --git a/source/tests/common/dpmodel/test_dist_check.py b/source/tests/common/dpmodel/test_dist_check.py index f64034e59d..7f8fb9b929 100644 --- a/source/tests/common/dpmodel/test_dist_check.py +++ b/source/tests/common/dpmodel/test_dist_check.py @@ -1,17 +1,27 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Unit tests for min_pair_dist frame filtering.""" +"""Unit tests for the pair-clearance frame filter.""" import unittest import numpy as np from deepmd.dpmodel.utils.dist_check import ( - compute_min_pair_dist_single, + compute_min_pair_margin_batch, + compute_min_pair_margin_single, + pair_half_thresholds, + select_frames, ) +from deepmd.utils.model_stat import ( + make_stat_input, +) + +#: A unit length scale: every pair threshold is 1 Å, so a margin is the pair +#: distance itself and the filter reduces to a plain distance floor. +UNIT_HALF = np.full(8, 0.5) -class TestComputeMinPairDistSingle(unittest.TestCase): - """Test minimum pairwise distance computation.""" +class TestComputeMinPairMarginSingle(unittest.TestCase): + """Test minimum pair margin computation on a unit length scale.""" def test_three_atoms_no_pbc(self) -> None: """Three atoms, closest pair is 0.3 Å.""" @@ -29,7 +39,7 @@ def test_three_atoms_no_pbc(self) -> None: ] ) atype = np.array([0, 0, 1]) - dist = compute_min_pair_dist_single(coord, box=None, atype=atype) + dist = compute_min_pair_margin_single(coord, None, atype, UNIT_HALF) np.testing.assert_almost_equal(dist, 0.3) def test_pbc_minimum_image(self) -> None: @@ -40,7 +50,7 @@ def test_pbc_minimum_image(self) -> None: coord = np.array([0.5, 5.0, 5.0, 9.5, 5.0, 5.0]) box = np.array([10.0, 0, 0, 0, 10.0, 0, 0, 0, 10.0]) atype = np.array([0, 0]) - dist = compute_min_pair_dist_single(coord, box=box, atype=atype) + dist = compute_min_pair_margin_single(coord, box, atype, UNIT_HALF) np.testing.assert_almost_equal(dist, 1.0) def test_pbc_triclinic(self) -> None: @@ -49,7 +59,7 @@ def test_pbc_triclinic(self) -> None: box = np.array([10.0, 0, 0, 2.0, 10.0, 0, 0, 0, 10.0]) coord = np.array([0.2, 0.0, 0.0, 9.8, 0.0, 0.0]) atype = np.array([0, 0]) - dist = compute_min_pair_dist_single(coord, box=box, atype=atype) + dist = compute_min_pair_margin_single(coord, box, atype, UNIT_HALF) np.testing.assert_almost_equal(dist, 0.4, decimal=5) def test_virtual_atoms_excluded(self) -> None: @@ -68,46 +78,42 @@ def test_virtual_atoms_excluded(self) -> None: ] ) atype = np.array([0, -1, 1]) - dist = compute_min_pair_dist_single(coord, box=None, atype=atype) + dist = compute_min_pair_margin_single(coord, None, atype, UNIT_HALF) np.testing.assert_almost_equal(dist, 2.0) def test_single_real_atom(self) -> None: """Only one real atom returns inf.""" coord = np.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]) atype = np.array([0, -1]) - dist = compute_min_pair_dist_single(coord, box=None, atype=atype) + dist = compute_min_pair_margin_single(coord, None, atype, UNIT_HALF) self.assertEqual(dist, float("inf")) def test_all_virtual(self) -> None: """All virtual atoms return inf.""" coord = np.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]) atype = np.array([-1, -1]) - dist = compute_min_pair_dist_single(coord, box=None, atype=atype) + dist = compute_min_pair_margin_single(coord, None, atype, UNIT_HALF) self.assertEqual(dist, float("inf")) def test_coord_shape_2d(self) -> None: """Accept (natoms, 3) shaped coord.""" coord = np.array([[0.0, 0.0, 0.0], [0.8, 0.0, 0.0]]) atype = np.array([0, 1]) - dist = compute_min_pair_dist_single(coord, box=None, atype=atype) + dist = compute_min_pair_margin_single(coord, None, atype, UNIT_HALF) np.testing.assert_almost_equal(dist, 0.8) - def test_stop_below_triggers_early_exit(self) -> None: - """A pair below stop_below should still return the correct minimum.""" + def test_a_pair_inside_its_window_owns_the_margin(self) -> None: + """A pair far inside its window is the one the scan reports.""" coord = np.array([0.0, 0.0, 0.0, 0.05, 0.0, 0.0, 10.0, 0.0, 0.0]) atype = np.array([0, 0, 0]) - dist = compute_min_pair_dist_single( - coord, box=None, atype=atype, stop_below=0.1 - ) + dist = compute_min_pair_margin_single(coord, None, atype, UNIT_HALF) np.testing.assert_almost_equal(dist, 0.05) - def test_stop_below_not_triggered(self) -> None: - """If all pairs are above stop_below, the true minimum is returned.""" + def test_a_pair_at_the_edge_of_its_window_reports_one(self) -> None: + """A frame whose closest pair sits exactly on its window reads one.""" coord = np.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 2.0, 0.0, 0.0]) atype = np.array([0, 0, 0]) - dist = compute_min_pair_dist_single( - coord, box=None, atype=atype, stop_below=0.5 - ) + dist = compute_min_pair_margin_single(coord, None, atype, UNIT_HALF) np.testing.assert_almost_equal(dist, 1.0) def test_multi_block_iteration(self) -> None: @@ -121,16 +127,496 @@ def test_multi_block_iteration(self) -> None: np.fill_diagonal(dist, np.inf) ref = dist.min() - actual = compute_min_pair_dist_single(coord, box=None, atype=atype) + actual = compute_min_pair_margin_single(coord, None, atype, UNIT_HALF) np.testing.assert_almost_equal(actual, ref, decimal=10) def test_coincident_atoms_zero(self) -> None: """Coincident real atoms should return exactly zero.""" coord = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0]) atype = np.array([0, 0, 0]) - dist = compute_min_pair_dist_single(coord, box=None, atype=atype) + dist = compute_min_pair_margin_single(coord, None, atype, UNIT_HALF) self.assertEqual(dist, 0.0) + def test_a_screened_scan_keeps_the_side_of_the_window(self) -> None: + """A screened scan stops at the first pair inside the window. + + It then reports the margin it had reached rather than the smallest one, + which can only be larger and still lies below one. A frame that clears + its window is never cut short, so its margin is the exact one. + """ + # A lattice wide enough for several row blocks, so a scan that stops + # early stops before reaching the closest pair. + grid = np.arange(9) * 2.0 + coord = np.stack(np.meshgrid(grid, grid, grid, indexing="ij"), axis=-1) + coord = coord.reshape(-1, 3) + atype = np.zeros(coord.shape[0], dtype=np.int64) + clear = compute_min_pair_margin_single(coord, None, atype, UNIT_HALF) + self.assertGreater(clear, 1.0) + self.assertEqual( + compute_min_pair_margin_single( + coord, None, atype, UNIT_HALF, screened=True + ), + clear, + ) + + # One pair inside the window early in the scan, a closer one late. + coord[1] = coord[0] + np.array([0.7, 0.0, 0.0]) + coord[721] = coord[720] + np.array([0.2, 0.0, 0.0]) + np.testing.assert_allclose( + compute_min_pair_margin_single(coord, None, atype, UNIT_HALF), + 0.2, + rtol=1e-12, + ) + np.testing.assert_allclose( + compute_min_pair_margin_single( + coord, None, atype, UNIT_HALF, screened=True + ), + 0.7, + rtol=1e-12, + ) + + +class TestComputeMinPairMarginBatch(unittest.TestCase): + """The batched scan gives every frame the margin of the single scan.""" + + @staticmethod + def _frames( + nframes: int, + natoms: int, + seed: int, + *, + periodic: bool = True, + virtual: bool = False, + clustered: bool = False, + sizes: np.ndarray | None = None, + ) -> tuple[list[np.ndarray], np.ndarray, list[np.ndarray]]: + """Build frames of random geometry, one cell each.""" + rng = np.random.default_rng(seed) + counts = np.full(nframes, natoms) if sizes is None else sizes + coords, cells, types = [], [], [] + for count in counts: + cell = np.diag([7.0, 8.5, 11.0]) + 0.6 * rng.standard_normal((3, 3)) + fractional = rng.uniform(0.0, 1.0, (count, 3)) + if clustered: + # Two tight clusters half a cell apart: the closest pair of a + # frame may cross the periodic boundary either way. + fractional[: count // 2] *= 0.05 + fractional[count // 2 :] = 0.55 + 0.05 * fractional[count // 2 :] + coord = fractional @ cell + if count > 1 and rng.random() < 0.5: + # Half the frames hold a pair far below any plausible bound. + coord[1] = coord[0] + 0.05 * rng.standard_normal(3) + atype = rng.integers(0, 4, count) + if virtual: + atype[rng.random(count) < 0.25] = -1 + coords.append(coord) + cells.append(cell.reshape(9) if periodic else None) + types.append(atype) + boxes = np.stack(cells) if periodic else None + return coords, boxes, types + + def _assert_matches_single( + self, coords, boxes, types, half=UNIT_HALF, **kwargs + ) -> None: + expected = np.array( + [ + compute_min_pair_margin_single( + coord, None if boxes is None else boxes[frame], atype, half + ) + for frame, (coord, atype) in enumerate(zip(coords, types, strict=True)) + ] + ) + got = compute_min_pair_margin_batch( + np.stack(coords) if "n_node" not in kwargs else np.concatenate(coords), + boxes, + np.stack(types) if "n_node" not in kwargs else np.concatenate(types), + half, + **kwargs, + ) + np.testing.assert_allclose(got, expected, rtol=1e-12, atol=1e-12) + + def test_matches_the_single_frame_scan(self) -> None: + for periodic in (True, False): + for virtual in (False, True): + for clustered in (False, True): + with self.subTest( + periodic=periodic, virtual=virtual, clustered=clustered + ): + self._assert_matches_single( + *self._frames( + 17, + 14, + seed=3, + periodic=periodic, + virtual=virtual, + clustered=clustered, + ) + ) + + def test_ragged_batch_matches_the_single_frame_scan(self) -> None: + sizes = np.array([5, 12, 2, 9, 12]) + coords, boxes, types = self._frames(len(sizes), 0, seed=5, sizes=sizes) + self._assert_matches_single(coords, boxes, types, n_node=sizes) + + def test_frames_without_a_pair_carry_infinity(self) -> None: + coords, boxes, types = self._frames(3, 1, seed=7) + got = compute_min_pair_margin_batch( + np.stack(coords), boxes, np.stack(types), UNIT_HALF + ) + np.testing.assert_array_equal(got, np.full(3, np.inf)) + + def test_periodic_and_open_frames_mix_in_one_batch(self) -> None: + """A cell of zeros marks a frame the minimum image does not apply to.""" + coords, boxes, types = self._frames(11, 9, seed=13) + boxes = boxes.copy() + boxes[[2, 5, 6]] = 0.0 + expected = np.array( + [ + compute_min_pair_margin_single( + coord, + boxes[frame] if boxes[frame].any() else None, + atype, + UNIT_HALF, + ) + for frame, (coord, atype) in enumerate(zip(coords, types, strict=True)) + ] + ) + got = compute_min_pair_margin_batch( + np.stack(coords), boxes, np.stack(types), UNIT_HALF + ) + np.testing.assert_allclose(got, expected, rtol=1e-12, atol=1e-12) + + def test_filler_atoms_leave_the_scan(self) -> None: + """Padding rows share one position and must not hold up the scan. + + A rectangular batch pads its frames with atoms of a negative type at + the origin. They carry no distance, and leaving them in the order + would stall the window, since they never separate. + """ + coords, boxes, types = self._frames(12, 7, seed=17) + padded_coords, padded_types = [], [] + for coord, atype in zip(coords, types, strict=True): + padded_coords.append(np.concatenate([coord, np.zeros((25, 3))])) + padded_types.append(np.concatenate([atype, np.full(25, -1)])) + expected = compute_min_pair_margin_batch( + np.stack(coords), boxes, np.stack(types), UNIT_HALF + ) + got = compute_min_pair_margin_batch( + np.stack(padded_coords), boxes, np.stack(padded_types), UNIT_HALF + ) + np.testing.assert_allclose(got, expected, rtol=1e-12, atol=1e-12) + + def test_the_exact_scan_reaches_beyond_the_first_neighbours(self) -> None: + """Without a bound the scan widens until it owns the minimum. + + The closest pair of every frame is hidden behind several atoms that + separate it along the ordered direction, so neither the first windows + nor a scan that screened at a distance of its own would find it. + """ + rng = np.random.default_rng(23) + nframes, natoms, edge = 9, 26, 14.0 + cells = np.stack([np.diag([edge, edge, edge]) for _ in range(nframes)]) + atype = np.zeros((nframes, natoms), dtype=np.int64) + coords = np.empty((nframes, natoms, 3)) + for frame in range(nframes): + # A row of atoms along x, spread far apart in y so that only the + # planted pair is close in space. + position = np.stack( + [ + np.linspace(1.0, 12.0, natoms), + rng.permutation(np.linspace(0.0, 12.0, natoms)), + np.zeros(natoms), + ], + axis=-1, + ) + # The winning pair straddles eight of those atoms. + position[3] = [4.0, 6.0, 0.0] + position[12] = [4.3, 6.0, 0.0] + coords[frame] = position + boxes = cells.reshape(nframes, 9) + expected = np.array( + [ + compute_min_pair_margin_single( + coords[frame], boxes[frame], atype[frame], UNIT_HALF + ) + for frame in range(nframes) + ] + ) + np.testing.assert_allclose(expected, 0.3, rtol=1e-12, atol=1e-12) + # The answer is not among the pairs that neighbour each other in the + # ordered direction, so the scan cannot stop at its first window. + order = np.argsort(coords[:, :, 0], axis=1) + ordered = np.take_along_axis(coords, order[:, :, np.newaxis], axis=1) + adjacent = np.linalg.norm(ordered[:, 1:] - ordered[:, :-1], axis=-1).min(axis=1) + self.assertTrue((adjacent > expected + 1e-9).all()) + got = compute_min_pair_margin_batch(coords, boxes, atype, UNIT_HALF) + np.testing.assert_allclose(got, expected, rtol=1e-12, atol=1e-12) + + def test_the_window_closes_on_the_distance_it_has_reached(self) -> None: + """Two periodic frames on which a loose stopping rule reports too much. + + The first frame ends the scan one window early if the separation the + window has travelled is taken twice as long as it is. The second needs + that separation measured in the direction of the order: reading it + through the periodic image instead lets the window close before the + pair its other atom would have met. + """ + edges = (6.7731, 7.9566) + positions = ( + [ + [0.8561, 5.8872, 4.4824], + [4.1117, 5.0969, 6.2866], + [2.0466, 0.3791, 4.5902], + [2.1129, 4.8734, 4.6289], + [4.5744, 5.5052, 1.0678], + [1.2241, 3.5232, 1.1324], + ], + [ + [7.6215, 2.5614, 1.8946], + [6.9715, 6.4405, 7.6013], + [2.2376, 1.8901, 1.6561], + [7.7785, 7.3701, 4.1730], + [2.2114, 4.5956, 4.2626], + [4.6901, 5.2300, 3.8674], + [0.3614, 6.0416, 7.4999], + [0.1976, 4.1174, 2.7984], + ], + ) + for edge, coord in zip(edges, positions, strict=True): + with self.subTest(edge=edge): + frame = np.array(coord) + box = np.diag([edge, edge, edge]).reshape(1, 9) + atype = np.zeros((1, frame.shape[0]), dtype=np.int64) + expected = compute_min_pair_margin_single( + frame, box[0], atype[0], UNIT_HALF + ) + got = compute_min_pair_margin_batch( + frame[np.newaxis], box, atype, UNIT_HALF + ) + np.testing.assert_allclose(got, [expected], rtol=1e-12, atol=1e-12) + + def test_a_screening_bound_keeps_the_side_of_every_frame(self) -> None: + """A bounded scan keeps the margin of every frame that fails. + + Bounding the search leaves a frame that clears its window with a + margin it reaches rather than its minimum, so the value may grow; a + frame that fails keeps the exact margin, since the bound is then no + longer the smaller of the two. + """ + for half in (UNIT_HALF, np.array([0.18, 0.31, 0.46, 0.75, 0.2, 0.2, 0.2, 0.2])): + with self.subTest(half=half[0]): + coords, boxes, types = self._frames(23, 16, seed=11) + exact = compute_min_pair_margin_batch( + np.stack(coords), boxes, np.stack(types), half + ) + screened = compute_min_pair_margin_batch( + np.stack(coords), boxes, np.stack(types), half, screened=True + ) + self.assertTrue((exact < 1.0).any() and (exact >= 1.0).any()) + np.testing.assert_array_equal(screened < 1.0, exact < 1.0) + np.testing.assert_allclose( + screened[exact < 1.0], exact[exact < 1.0], rtol=1e-12, atol=1e-12 + ) + self.assertTrue((screened >= exact - 1e-12).all()) + + +class TestPairLengthScale(unittest.TestCase): + """The window each pair carries is the sum of its two half-thresholds.""" + + def test_absolute_scale_is_a_plain_distance_floor(self) -> None: + """Every element gets the unit radius. + + A pair's threshold is then the fraction itself and the margin is the + pair distance divided by it. + """ + half = pair_half_thresholds(0.5, "absolute", ntypes=3) + np.testing.assert_allclose(half, 0.25) + coord = np.array([0.0, 0.0, 0.0, 0.8, 0.0, 0.0]) + atype = np.array([0, 2]) + margin = compute_min_pair_margin_single(coord, None, atype, half) + np.testing.assert_allclose(margin, 0.8 / 0.5) + + def test_covalent_scale_sizes_each_pair_separately(self) -> None: + """Two pairs at one distance carry different margins.""" + half = np.array([0.1, 0.4]) + coord = np.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]) + close = compute_min_pair_margin_single(coord, None, np.array([1, 1]), half) + wide = compute_min_pair_margin_single(coord, None, np.array([0, 0]), half) + np.testing.assert_allclose(close, 1.0 / 0.8) + np.testing.assert_allclose(wide, 1.0 / 0.2) + + def test_the_minimum_margin_pair_is_not_the_minimum_distance_pair(self) -> None: + """The scan tracks the margin, which a closer pair need not own. + + The last two atoms are the closest pair of the frame, but they are + small and their window is narrow, so they clear it comfortably. The + wider-apart pair of large atoms sits deeper inside its own window and + owns the margin. + """ + half = np.array([0.05, 1.0]) + coord = np.array([0.0, 0.0, 0.0, 1.5, 0.0, 0.0, 6.0, 0.0, 0.0, 6.4, 0.0, 0.0]) + atype = np.array([1, 1, 0, 0]) + distances = np.array([1.5, 0.4]) + margins = distances / np.array([2.0, 0.1]) + self.assertLess(distances[1], distances[0]) + self.assertLess(margins[0], margins[1]) + got = compute_min_pair_margin_single(coord, None, atype, half) + np.testing.assert_allclose(got, margins[0]) + batched = compute_min_pair_margin_batch( + coord[np.newaxis], None, atype[np.newaxis], half + ) + np.testing.assert_allclose(batched, [margins[0]]) + + def test_batch_matches_single_under_a_per_pair_scale(self) -> None: + """The batched scan owns the margin for element-dependent windows.""" + half = np.array([0.18, 0.31, 0.46, 0.75, 0.2, 0.2, 0.2, 0.2]) + for clustered in (False, True): + with self.subTest(clustered=clustered): + coords, boxes, types = TestComputeMinPairMarginBatch._frames( + 19, 15, seed=29, clustered=clustered + ) + expected = np.array( + [ + compute_min_pair_margin_single(coord, boxes[frame], atype, half) + for frame, (coord, atype) in enumerate( + zip(coords, types, strict=True) + ) + ] + ) + got = compute_min_pair_margin_batch( + np.stack(coords), boxes, np.stack(types), half + ) + np.testing.assert_allclose(got, expected, rtol=1e-12, atol=1e-12) + + def test_covalent_scale_needs_a_type_map(self) -> None: + with self.assertRaisesRegex(ValueError, "type_map"): + pair_half_thresholds(0.26, "covalent", ntypes=3) + + def test_covalent_scale_reads_the_periodic_table(self) -> None: + half = pair_half_thresholds(0.5, "covalent", type_map=["H", "O"]) + # Half of the Pyykko single-bond radii of hydrogen and oxygen. + np.testing.assert_allclose(half, [0.16, 0.315]) + + def test_unknown_scale_is_refused(self) -> None: + with self.assertRaisesRegex(ValueError, "length scale"): + pair_half_thresholds(0.26, "vdw", ntypes=3) + + +class TestSelectFrames(unittest.TestCase): + """`select_frames` slices the frame axis and, when ragged, the node axis.""" + + def test_rectangular_batch(self) -> None: + batch = { + "coord": np.arange(3 * 6, dtype=np.float64).reshape(3, 6), + "energy": np.array([[1.0], [2.0], [3.0]]), + "find_energy": np.bool_(True), + "box": None, + } + kept = select_frames(batch, np.array([True, False, True])) + np.testing.assert_array_equal(kept["coord"], batch["coord"][[0, 2]]) + np.testing.assert_array_equal(kept["energy"][:, 0], [1.0, 3.0]) + self.assertTrue(kept["find_energy"]) + self.assertIsNone(kept["box"]) + + def test_ragged_batch(self) -> None: + # Three frames of two, one and three atoms on one flat node axis. + n_node = np.array([2, 1, 3]) + frame_of_node = np.repeat(np.arange(3), n_node) + batch = { + "n_node": n_node, + "atype": frame_of_node.copy(), + "coord": np.repeat(frame_of_node, 3).reshape(6, 3).astype(np.float64), + # A per-atom field stored flat, with two entries per atom. + "aparam": np.repeat(frame_of_node, 2).astype(np.float64), + "energy": np.array([[10.0], [11.0], [12.0]]), + "box": np.arange(27, dtype=np.float64).reshape(3, 9), + } + kept = select_frames(batch, np.array([True, False, True])) + np.testing.assert_array_equal(kept["n_node"], [2, 3]) + np.testing.assert_array_equal(kept["atype"], [0, 0, 2, 2, 2]) + np.testing.assert_array_equal(kept["coord"][:, 0], [0, 0, 2, 2, 2]) + np.testing.assert_array_equal(kept["aparam"], np.repeat([0, 0, 2, 2, 2], 2)) + np.testing.assert_array_equal(kept["energy"][:, 0], [10.0, 12.0]) + np.testing.assert_array_equal(kept["box"], batch["box"][[0, 2]]) + + def test_ragged_batch_of_single_atoms(self) -> None: + # One atom per frame makes the frame axis and the node axis coincide. + batch = { + "n_node": np.ones(3, dtype=np.int64), + "atype": np.array([0, 1, 2]), + "energy": np.array([[10.0], [11.0], [12.0]]), + } + kept = select_frames(batch, np.array([False, True, True])) + np.testing.assert_array_equal(kept["atype"], [1, 2]) + np.testing.assert_array_equal(kept["energy"][:, 0], [11.0, 12.0]) + np.testing.assert_array_equal(kept["n_node"], [1, 1]) + + +class _CyclicDataSystem: + """Data source of one-frame batches that cycle through each system. + + Frame ``k`` of a system is tagged by ``coord == k`` and carries the listed + pair margin, so a test reads off which frames were kept. + """ + + def __init__(self, margins: list[list[float]], derived: bool = True) -> None: + self.margins = margins + self.derived = derived + self.cursor = [0] * len(margins) + + def get_nsystems(self) -> int: + return len(self.margins) + + def get_nbatches(self) -> list[int]: + return [len(system) for system in self.margins] + + def get_batch(self, sys_idx: int) -> dict: + frame = self.cursor[sys_idx] % len(self.margins[sys_idx]) + self.cursor[sys_idx] += 1 + return { + "coord": np.full((1, 6), float(frame)), + "type": np.zeros((1, 2), dtype=np.int32), + "natoms_vec": np.array([2, 2, 2]), + "pair_margin": np.array([[self.margins[sys_idx][frame]]]), + "find_pair_margin": float(self.derived), + } + + +class TestStatisticsFilter(unittest.TestCase): + """`make_stat_input` packs the frames whose pairs clear their window.""" + + def test_scan_replaces_filtered_batches(self) -> None: + data = _CyclicDataSystem([[0.2, 0.3, 1.5, 0.1, 1.2], [0.1, 0.2, 0.3]]) + packed = make_stat_input(data, nbatches=2, min_pair_dist=1.0) + # The second system holds no valid frame and is skipped after one + # pass; the first one yields its two valid frames. + self.assertEqual(len(packed), 1) + np.testing.assert_array_equal(packed[0]["coord"][:, 0], [2.0, 4.0]) + np.testing.assert_array_equal(packed[0]["pair_margin"][:, 0], [1.5, 1.2]) + self.assertEqual(data.cursor, [5, 3]) + + def test_scan_stops_at_the_requested_batches(self) -> None: + data = _CyclicDataSystem([[1.5, 0.2, 1.2, 1.4]]) + packed = make_stat_input(data, nbatches=2, min_pair_dist=1.0) + np.testing.assert_array_equal(packed[0]["coord"][:, 0], [0.0, 2.0]) + self.assertEqual(data.cursor, [3]) + + def test_disabled_filter_takes_the_next_batches(self) -> None: + data = _CyclicDataSystem([[0.2, 0.3, 1.5], [0.1, 0.2]]) + packed = make_stat_input(data, nbatches=3) + np.testing.assert_array_equal(packed[0]["coord"][:, 0], [0.0, 1.0, 2.0]) + np.testing.assert_array_equal(packed[1]["coord"][:, 0], [0.0, 1.0, 0.0]) + + def test_dataset_without_valid_frames_is_reported(self) -> None: + data = _CyclicDataSystem([[0.2, 0.3], [0.1]]) + with self.assertRaisesRegex(RuntimeError, "beyond the filter radius"): + make_stat_input(data, nbatches=2, min_pair_dist=1.0) + + def test_underived_distance_is_refused(self) -> None: + data = _CyclicDataSystem([[1.5, 1.2]], derived=False) + with self.assertRaisesRegex(RuntimeError, "did not derive"): + make_stat_input(data, nbatches=1, min_pair_dist=1.0) + if __name__ == "__main__": unittest.main() diff --git a/source/tests/common/dpmodel/test_dpa4_call_graph.py b/source/tests/common/dpmodel/test_dpa4_call_graph.py index c2410f2e5c..44c65bfe19 100644 --- a/source/tests/common/dpmodel/test_dpa4_call_graph.py +++ b/source/tests/common/dpmodel/test_dpa4_call_graph.py @@ -147,8 +147,9 @@ def _run_graph(dd, coord, atype, nlist, permute_seed=None): edge_vec=graph.edge_vec[perm], edge_mask=graph.edge_mask[perm], ) - out, rot_mat = dd.call_graph(graph, atype.reshape(-1)) + out, rot_mat, node_gate = dd.call_graph(graph, atype.reshape(-1)) assert rot_mat is None + assert node_gate is None return np.asarray(out) @@ -202,8 +203,8 @@ def test_call_graph_spin_sensitivity() -> None: 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) + 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)) @@ -218,7 +219,7 @@ def test_call_graph_spin_matches_dense_adapter() -> None: 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) + out, _, _ = dd.call_graph(graph, atype_flat, spin=spin) np.testing.assert_allclose( np.asarray(out).reshape(nf, nloc, -1), np.asarray(ref), @@ -267,8 +268,8 @@ def test_call_graph_charge_spin_sensitivity() -> None: 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) + 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)) @@ -286,7 +287,7 @@ def test_call_graph_charge_spin_matches_dense_adapter() -> None: 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) + out, _, _ = dd.call_graph(graph, atype_flat, charge_spin=cs) np.testing.assert_allclose( np.asarray(out).reshape(nf, nloc, -1), np.asarray(ref), @@ -352,8 +353,8 @@ def test_message_sensitive_fixture_is_edge_dependent() -> None: 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, _, _ = 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 @@ -496,11 +497,12 @@ def test_output_stat_graph_wrapper_flattens_aparam() -> None: 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 + ``inner_clamp_f_inner``/``inner_clamp_f_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``. + same construction. The absolute scale gives every pair the unit length + scale, so the two fractions are the window radii in Å. Jittered for the + same message-sensitivity reason as ``make_message_sensitive_descriptor``. """ dd = DescrptDPA4( ntypes=2, @@ -514,8 +516,9 @@ def make_bridging_descriptor(seed: int = 99) -> DescrptDPA4: precision="float64", seed=7, random_gamma=False, - inner_clamp_r_inner=0.8, - inner_clamp_r_outer=1.2, + inner_clamp_f_inner=0.8, + inner_clamp_f_outer=1.2, + inner_clamp_scale="absolute", ) data = dd.serialize() data = jitter_zero_arrays(data, np.random.default_rng(seed)) @@ -537,17 +540,21 @@ def _sphere_points(n_points: int, radius: float) -> np.ndarray: def _frozen_sphere_descriptor_outputs( dd: DescrptDPA4, near_distance: float, n_points: int = 12 -) -> np.ndarray: +) -> tuple[np.ndarray, np.ndarray | None]: """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). + both. Returns the flat call_graph descriptor reshaped to + ``(n_points, 3, channels)`` and the per-node source gate reshaped to + ``(n_points, 3)`` (``None`` for an unbridged descriptor). 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). + ``_evaluate_frozen_sphere_atom_energies`` at the descriptor level. The + gate leaves the two edges of the frozen pair alive, so the descriptors of + A and B keep following each other; what leaves the graph are the four + edges the frozen pair emits into C, which is where the invariance lands. + An invariant descriptor of C carries over to any per-atom fitting net, + itself a function of C's descriptor and type alone. """ directions = _sphere_points(n_points, near_distance) coord = np.zeros((n_points, 3, 3), dtype=np.float64) @@ -556,18 +563,31 @@ def _frozen_sphere_descriptor_outputs( 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) + out, _, node_gate = dd.call_graph(graph, atype.reshape(-1)) + return ( + np.asarray(out).reshape(n_points, 3, -1), + None if node_gate is None else np.asarray(node_gate).reshape(n_points, 3), + ) 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). + """SFPG bridging on the graph route: the frozen pair (A, B) must vanish + from the environment of the third atom C, so C's descriptor is invariant + to their rigid motion inside the frozen zone (below r_inner=0.8). + + The pair itself stays connected -- the leave-one-out gate keeps the two + edges between A and B -- so A's descriptor still follows the direction of + B; what makes that harmless is the per-node gate, which is exactly zero on + both frozen atoms and one on C. """ dd = make_bridging_descriptor() - out = _frozen_sphere_descriptor_outputs(dd, near_distance=0.5) + out, node_gate = _frozen_sphere_descriptor_outputs(dd, near_distance=0.5) + span_c = np.max(np.abs(out[:, 2, :] - out[0:1, 2, :])) + assert span_c < 1e-10, span_c span_a = np.max(np.abs(out[:, 0, :] - out[0:1, 0, :])) - assert span_a < 1e-10, span_a + assert span_a > 1e-6, span_a + np.testing.assert_array_equal(node_gate[:, :2], np.zeros((out.shape[0], 2))) + np.testing.assert_array_equal(node_gate[:, 2], np.ones(out.shape[0])) def test_call_graph_bridging_leak_reopens_when_gate_disabled() -> None: @@ -577,9 +597,10 @@ def test_call_graph_bridging_leak_reopens_when_gate_disabled() -> None: """ 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 + out, node_gate = _frozen_sphere_descriptor_outputs(dd, near_distance=0.5) + assert node_gate is None + span_c = np.max(np.abs(out[:, 2, :] - out[0:1, 2, :])) + assert span_c > 1e-6, span_c def test_charge_spin_model_routes_through_graph_lower() -> None: diff --git a/source/tests/common/dpmodel/test_dpa4_edge_cache.py b/source/tests/common/dpmodel/test_dpa4_edge_cache.py index 17f823f69f..263a597420 100644 --- a/source/tests/common/dpmodel/test_dpa4_edge_cache.py +++ b/source/tests/common/dpmodel/test_dpa4_edge_cache.py @@ -38,8 +38,9 @@ def fused_wigner(quaternion: np.ndarray) -> tuple[np.ndarray, np.ndarray]: compute_dtype=np.float64, eps=1.0e-12, deg_norm_floor=1.0, - inner_clamp=None, + bridging_clamp=None, bridging_switch=None, + edge_contact=None, edge_envelope=unexpected_reference, radial_basis=unexpected_reference, random_gamma=False, diff --git a/source/tests/common/dpmodel/test_dpa4_so2_grid.py b/source/tests/common/dpmodel/test_dpa4_so2_grid.py index 64934bfbd1..75ff318a6d 100644 --- a/source/tests/common/dpmodel/test_dpa4_so2_grid.py +++ b/source/tests/common/dpmodel/test_dpa4_so2_grid.py @@ -158,7 +158,6 @@ def _build_edge_data( inv_sqrt_deg=t(inv_sqrt_deg), D_full=t(D_full[valid]), Dt_full=t(Dt_full[valid]), - edge_src_gate=None, ) dp_cache = EdgeCache( src=src, @@ -171,7 +170,6 @@ def _build_edge_data( inv_sqrt_deg=inv_sqrt_deg, D_full=D_full, Dt_full=Dt_full, - edge_src_gate=None, edge_mask=mask, ) return pt_cache, dp_cache, radial, radial[valid], x diff --git a/source/tests/common/dpmodel/test_fitting_node_gate.py b/source/tests/common/dpmodel/test_fitting_node_gate.py new file mode 100644 index 0000000000..a593739f6a --- /dev/null +++ b/source/tests/common/dpmodel/test_fitting_node_gate.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""The readout gate of a bridging window on the dpmodel fitting. + +A bridged descriptor hands the fitting a per-atom source gate. The gate fades +the learned part of the output, the deviation of the output from the per-type +bias the fitting stores, so an atom whose gate is zero keeps that bias alone +and an atom whose gate is one is untouched. The same rule applies on the dense +call and on the graph-native one. +""" + +import numpy as np +import pytest + +from deepmd.dpmodel.fitting import ( + InvarFitting, +) + +NF, NLOC, ND, NTYPES = 2, 4, 8, 2 + + +def _fitting(vacuum_ref: bool) -> InvarFitting: + return InvarFitting( + "energy", + NTYPES, + ND, + 1, + mixed_types=True, + vacuum_ref=vacuum_ref, + seed=3, + ) + + +def _inputs(rng: np.random.Generator, vacuum_ref: bool): + descriptor = rng.normal(size=(NF, NLOC, ND)) + atype = rng.integers(0, NTYPES, size=(NF, NLOC)) + node_gate = rng.uniform(0.0, 1.0, size=(NF, NLOC, 1)) + vacuum = rng.normal(size=(NTYPES, ND)) if vacuum_ref else None + return descriptor, atype, node_gate, vacuum + + +@pytest.mark.parametrize("vacuum_ref", [False, True]) +def test_dense_call_fades_the_learned_part(vacuum_ref: bool) -> None: + """``out = bias + gate * (out_without_gate - bias)`` atom by atom.""" + rng = np.random.default_rng(0) + fitting = _fitting(vacuum_ref) + bias = np.asarray(fitting.bias_atom_e).reshape(NTYPES, 1) + descriptor, atype, node_gate, vacuum = _inputs(rng, vacuum_ref) + + ungated = fitting(descriptor, atype, vacuum_descriptor=vacuum)["energy"] + gated = fitting(descriptor, atype, vacuum_descriptor=vacuum, node_gate=node_gate)[ + "energy" + ] + + reference = bias[atype] + np.testing.assert_allclose( + gated, reference + node_gate * (ungated - reference), rtol=0.0, atol=1e-12 + ) + # the fixture must separate the two: a gate of one would make it vacuous + assert np.abs(gated - ungated).max() > 1e-6 + + +@pytest.mark.parametrize("vacuum_ref", [False, True]) +def test_a_closed_gate_leaves_the_bias_and_an_open_one_changes_nothing( + vacuum_ref: bool, +) -> None: + """The two ends of the fade are the bias itself and the plain output.""" + rng = np.random.default_rng(1) + fitting = _fitting(vacuum_ref) + bias = np.asarray(fitting.bias_atom_e).reshape(NTYPES, 1) + descriptor, atype, _, vacuum = _inputs(rng, vacuum_ref) + + closed = fitting( + descriptor, + atype, + vacuum_descriptor=vacuum, + node_gate=np.zeros((NF, NLOC, 1)), + )["energy"] + np.testing.assert_allclose(closed, bias[atype], rtol=0.0, atol=1e-12) + + ungated = fitting(descriptor, atype, vacuum_descriptor=vacuum)["energy"] + opened = fitting( + descriptor, + atype, + vacuum_descriptor=vacuum, + node_gate=np.ones((NF, NLOC, 1)), + )["energy"] + np.testing.assert_allclose(opened, ungated, rtol=0.0, atol=1e-12) + + if vacuum_ref: + # the fold moves the reference out of the bias; the gate still reads + # the bias the atoms of a frozen pair fade to + fitting.fold_vacuum_reference(vacuum) + assert fitting.folded_reference is not None + np.testing.assert_allclose(fitting.readout_reference(), bias, atol=1e-12) + folded = fitting(descriptor, atype, node_gate=np.zeros((NF, NLOC, 1))) + np.testing.assert_allclose(folded["energy"], closed, rtol=0.0, atol=1e-12) + + +@pytest.mark.parametrize("vacuum_ref", [False, True]) +def test_call_graph_matches_the_dense_gate(vacuum_ref: bool) -> None: + """The flat node axis carries the gate the same way the dense call does.""" + rng = np.random.default_rng(2) + fitting = _fitting(vacuum_ref) + descriptor, atype, node_gate, vacuum = _inputs(rng, vacuum_ref) + + dense = fitting(descriptor, atype, vacuum_descriptor=vacuum, node_gate=node_gate)[ + "energy" + ] + n_nodes = NF * NLOC + flat = fitting.call_graph( + descriptor.reshape(n_nodes, ND), + atype.reshape(n_nodes), + vacuum_descriptor=vacuum, + node_gate=node_gate.reshape(n_nodes), + )["energy"] + assert flat.shape == (n_nodes, 1) + np.testing.assert_allclose(flat, dense.reshape(n_nodes, 1), rtol=0.0, atol=1e-12) diff --git a/source/tests/common/dpmodel/test_inner_potential.py b/source/tests/common/dpmodel/test_inner_potential.py index 87e0f6ab52..3c3679aa0c 100644 --- a/source/tests/common/dpmodel/test_inner_potential.py +++ b/source/tests/common/dpmodel/test_inner_potential.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""dpmodel ``InnerPotential`` (analytical ZBL bridging term) unit tests. +"""dpmodel ``InnerPotential`` (analytical bridging term) unit tests. Ports the pt reference values from ``source/tests/pt/model/test_sezm_model.py::TestInnerPotential``: the exact @@ -13,7 +13,9 @@ import pytest from deepmd.dpmodel.atomic_model.inner_potential import ( + ELEMENT_TO_Z, InnerPotential, + nlh_provenance, ) _A_BOHR = 0.5291772109 @@ -64,24 +66,6 @@ def test_zbl_known_value(type_map, atypes, zi, zj): ) -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 = InnerPotential(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 = InnerPotential(type_map=["O"]) edge_vec, edge_index, _ = _two_atom_inputs(0.8) @@ -114,15 +98,150 @@ def test_torch_namespace_smoke_and_gradient(): 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) + # The device is explicit: the pt test package, which other common tests + # import, points the default device at a nonexistent one on purpose. + ev = torch.tensor( + edge_vec_np, dtype=torch.float64, device="cpu", requires_grad=True + ) out = pot.call( ev, - torch.tensor(edge_index), - torch.tensor(atypes), - torch.tensor(edge_mask), + torch.tensor(edge_index, device="cpu"), + torch.tensor(atypes, device="cpu"), + torch.tensor(edge_mask, device="cpu"), 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 + + +# --------------------------------------------------------------------------- +# NLH mode: the bundled coefficient table and the conditions it guarantees +# --------------------------------------------------------------------------- +_ALL_ELEMENTS = list(ELEMENT_TO_Z) + + +def _nlh_series(pot: InnerPotential, ia: int, ib: int) -> tuple[np.ndarray, np.ndarray]: + """Amplitudes in eV Å and decay rates in Å⁻¹ of one ordered type pair.""" + stride = pot.ntypes + 1 + row = np.asarray(pot.series_table)[ia * stride + ib] + return row[:4], row[4:] + + +def test_nlh_mode_is_accepted_and_normalized(): + pot = InnerPotential(type_map=["O", "H"], mode="nlh") + assert pot.mode == "NLH" + + +def test_nlh_matches_its_table(): + """The scattered energy is the series the table holds for that pair.""" + r = 0.8 + pot = InnerPotential(type_map=["O", "H"], mode="nlh") + edge_vec, edge_index, edge_mask = _two_atom_inputs(r) + out = pot.call(edge_vec, edge_index, np.array([0, 1], dtype=np.int64), edge_mask, 2) + amp, rate = _nlh_series(pot, 0, 1) + expected = float(np.sum(amp * np.exp(-rate * r)) / r) + np.testing.assert_allclose(float(np.sum(out)), expected, rtol=1e-12) + + +def test_nlh_differs_from_zbl(): + """The two modes are distinct potentials, not the same table twice.""" + r = 0.8 + edge_vec, edge_index, edge_mask = _two_atom_inputs(r) + atypes = np.zeros(2, dtype=np.int64) + values = [ + float( + np.sum( + InnerPotential(type_map=["O"], mode=mode).call( + edge_vec, edge_index, atypes, edge_mask, 2 + ) + ) + ) + for mode in ("zbl", "nlh") + ] + assert values[1] > 0.0 + assert abs(values[1] / values[0] - 1.0) > 0.05 + + +def test_nlh_screening_reaches_the_coulomb_limit(): + """Every amplitude is non-negative and the amplitudes sum to k_e Z_a Z_b.""" + pot = InnerPotential(type_map=_ALL_ELEMENTS, mode="nlh") + charge = np.asarray([ELEMENT_TO_Z[symbol] for symbol in _ALL_ELEMENTS], dtype=float) + ntypes = len(_ALL_ELEMENTS) + table = np.asarray(pot.series_table).reshape(ntypes + 1, ntypes + 1, 8) + amp = table[:ntypes, :ntypes, :4] + assert amp.min() >= 0.0 + reference = _KE * charge[:, None] * charge[None, :] + np.testing.assert_allclose(amp.sum(-1), reference, rtol=1e-12) + + +def test_nlh_is_positive_and_strictly_decreasing(): + """Non-negative amplitudes make V positive and falling at every radius.""" + pot = InnerPotential(type_map=_ALL_ELEMENTS, mode="nlh") + ntypes = len(_ALL_ELEMENTS) + table = np.asarray(pot.series_table).reshape(ntypes + 1, ntypes + 1, 8) + radii = np.concatenate([np.geomspace(1e-3, 2.0, 200), np.linspace(2.02, 20.0, 200)]) + amp = table[:ntypes, :ntypes, None, :4] + rate = table[:ntypes, :ntypes, None, 4:] + v = (amp * np.exp(-rate * radii[:, None])).sum(-1) / radii + assert v.min() > 0.0 + assert np.diff(v, axis=-1).max() < 0.0 + + +def test_nlh_tail_bound_holds_for_every_pair(): + """The bound the fit imposes: 1 meV at 5 Å and 0.1 meV at 6 Å.""" + pot = InnerPotential(type_map=_ALL_ELEMENTS, mode="nlh") + ntypes = len(_ALL_ELEMENTS) + table = np.asarray(pot.series_table).reshape(ntypes + 1, ntypes + 1, 8) + amp, rate = table[:ntypes, :ntypes, :4], table[:ntypes, :ntypes, 4:] + for radius, bound in ((5.0, 1e-3), (6.0, 1e-4)): + v = (amp * np.exp(-rate * radius)).sum(-1) / radius + assert v.max() <= bound * (1.0 + 1e-6) + + +def test_nlh_table_is_symmetric_and_padded(): + """Swapping the two types leaves the row unchanged; padding rows vanish.""" + pot = InnerPotential(type_map=["H", "O", "U"], mode="nlh") + table = np.asarray(pot.series_table).reshape(4, 4, 8) + np.testing.assert_array_equal(table, np.swapaxes(table, 0, 1)) + np.testing.assert_array_equal(table[3], 0.0) + np.testing.assert_array_equal(table[:, 3], 0.0) + + +def test_nlh_change_type_map_rebuilds_from_symbols(): + """Reordering, adding and dropping elements all rebuild the right rows.""" + pot = InnerPotential(type_map=["O", "H"], mode="nlh") + original = np.asarray(_nlh_series(pot, 0, 1)[0]).copy() + pot.change_type_map(["H", "O", "Fe"]) + assert pot.type_map == ["H", "O", "Fe"] + assert pot.ntypes == 3 + assert np.asarray(pot.series_table).shape == (16, 8) + np.testing.assert_allclose(_nlh_series(pot, 1, 0)[0], original, rtol=1e-12) + fresh = InnerPotential(type_map=["H", "O", "Fe"], mode="nlh") + np.testing.assert_allclose( + np.asarray(pot.series_table), np.asarray(fresh.series_table), rtol=1e-12 + ) + + +def test_pair_table_is_the_float32_view_of_the_series_table(): + """The fused kernels and this route read the same constants.""" + for mode in ("zbl", "nlh"): + pot = InnerPotential(type_map=["O", "H"], mode=mode) + assert np.asarray(pot.pair_table).dtype == np.float32 + np.testing.assert_allclose( + np.asarray(pot.pair_table).astype(np.float64), + np.asarray(pot.series_table), + rtol=1e-6, + ) + + +def test_nlh_provenance_records_the_source_and_the_modification(): + """The attribution travels inside the data file.""" + meta = nlh_provenance() + assert "10.5281/zenodo.14172633" in meta["source_data"]["doi"] + assert "CC BY 4.0" in meta["source_data"]["licence"] + assert "10.1103/PhysRevA.111.032818" in meta["source_paper"]["doi"] + assert "not the coefficients published" in meta["modification"] + assert "MP2" in meta["rows"]["Z1 = 5, Z2 = 10"] + assert "ZBL" in meta["rows"]["Z1 > 92 or Z2 > 92"] diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py index 16dd1da1d1..304a4ba506 100644 --- a/source/tests/common/dpmodel/test_lmdb_data.py +++ b/source/tests/common/dpmodel/test_lmdb_data.py @@ -33,6 +33,10 @@ import numpy as np from deepmd.dpmodel.utils import lmdb_data as lmdb_data_module +from deepmd.dpmodel.utils.dist_check import ( + compute_min_pair_margin_single, + pair_half_thresholds, +) from deepmd.dpmodel.utils.lmdb_data import ( DistributedLmdbBatchSampler, LmdbBatchIterator, @@ -1202,17 +1206,20 @@ def test_lmdb_test_data(self): self.assertEqual(result["find_energy"], 1.0) self.assertEqual(result["find_force"], 1.0) - def test_min_pair_dist_requirement_computed(self): - path = _create_grid_lmdb(f"{self._tmpdir.name}/grid_min_pair.lmdb", nframes=1) + def test_pair_margin_requirement_computed(self): + path = _create_grid_lmdb( + f"{self._tmpdir.name}/grid_pair_margin.lmdb", nframes=1 + ) reader = LmdbDataReader(path, ["TYPE"], batch_size=1) reader.add_data_requirement( [ DataRequirementItem( - "min_pair_dist", + "pair_margin", ndof=1, atomic=False, must=False, high_prec=False, + default=0.5, source_policy="derived", ) ] @@ -1220,14 +1227,15 @@ def test_min_pair_dist_requirement_computed(self): frame = reader[0] - self.assertEqual(frame["find_min_pair_dist"], np.float32(1.0)) - np.testing.assert_allclose(frame["min_pair_dist"], np.array([1.0])) + # The grid spaces its atoms 1 Å apart, which is twice the window. + self.assertEqual(frame["find_pair_margin"], np.float32(1.0)) + np.testing.assert_allclose(frame["pair_margin"], np.array([2.0])) - def test_min_pair_dist_requirement_defaults_without_atype(self): + def test_pair_margin_requirement_defaults_without_atype(self): raw_frame = _make_frame(natoms=6, seed=0) raw_frame.pop("atom_types") requirement = DataRequirementItem( - "min_pair_dist", + "pair_margin", ndof=1, default=0.25, source_policy="derived", @@ -1236,7 +1244,7 @@ def test_min_pair_dist_requirement_defaults_without_atype(self): ntypes=2, natoms=6, type_remap=None, - data_requirements={"min_pair_dist": requirement}, + data_requirements={"pair_margin": requirement}, ) frame = decode_lmdb_frame( @@ -1246,20 +1254,20 @@ def test_min_pair_dist_requirement_defaults_without_atype(self): copy_arrays=True, ) - self.assertEqual(frame["find_min_pair_dist"], np.float32(0.0)) - np.testing.assert_allclose(frame["min_pair_dist"], np.array([0.25])) + self.assertEqual(frame["find_pair_margin"], np.float32(0.0)) + np.testing.assert_allclose(frame["pair_margin"], np.array([0.25])) - def test_derived_min_pair_dist_ignores_raw_presence_for_grouping(self): + def test_derived_pair_margin_ignores_raw_presence_for_grouping(self): """Derived fields neither partition frames nor trust stored values.""" path = _create_grid_lmdb( - f"{self._tmpdir.name}/derived_min_pair.lmdb", + f"{self._tmpdir.name}/derived_pair_margin.lmdb", nframes=2, ) environment = lmdb.open(path, readonly=False, lock=False) with environment.begin(write=True) as transaction: key = format(0, "012d").encode() frame = msgpack.unpackb(transaction.get(key), raw=False) - frame["min_pair_dist"] = { + frame["pair_margin"] = { "type": " None: self.assertFalse(iterator._pool.healthy) +class TestDerivedMinPairDist(unittest.TestCase): + """A batch derives `pair_margin` as the frames of that batch carry it.""" + + @staticmethod + def _requirement( + fraction: float, length_scale: str = "absolute" + ) -> list[DataRequirementItem]: + return [ + DataRequirementItem( + "pair_margin", + ndof=1, + atomic=False, + must=False, + high_prec=False, + default=fraction, + source_policy="derived", + length_scale=length_scale, + ) + ] + + def _margins(self, fraction: float, length_scale: str = "absolute") -> np.ndarray: + reader = LmdbDataReader(self._path, ["O", "H"], batch_size=6) + reader.add_data_requirement(self._requirement(fraction, length_scale)) + return reader.decode_batch(list(range(6)))["pair_margin"].reshape(-1) + + def _frame_margins( + self, fraction: float, length_scale: str = "absolute" + ) -> np.ndarray: + """Margins of the frame read path, which visits every pair. + + The batch read bounds its scan at a margin of one, so a frame that + clears its window carries a value it reaches rather than its exact + margin; the frame read path has no such bound. + """ + reader = LmdbDataReader(self._path, ["O", "H"], batch_size=6) + reader.add_data_requirement(self._requirement(fraction, length_scale)) + return np.concatenate([reader[index]["pair_margin"] for index in range(6)]) + + def setUp(self) -> None: + self._tmpdir = tempfile.TemporaryDirectory() + self._path = _create_lmdb(f"{self._tmpdir.name}/derived.lmdb", nframes=6) + # Frames without periodicity carry a cell of zeros, which the readers + # may mix with periodic ones in one batch. + environment = lmdb.open(self._path, map_size=10 * 1024 * 1024) + with environment.begin(write=True) as transaction: + for index in (1, 4): + key = format(index, "012d").encode() + frame = msgpack.unpackb(transaction.get(key), raw=False) + frame["cells"]["data"] = np.zeros((3, 3)).tobytes() + transaction.put(key, msgpack.packb(frame, use_bin_type=True)) + environment.close() + + def tearDown(self) -> None: + self._tmpdir.cleanup() + + def test_batch_matches_the_frames_it_holds(self) -> None: + """A window no frame clears leaves both read-backs exact.""" + wide = 10.0 + reader = LmdbDataReader(self._path, ["O", "H"], batch_size=6) + reader.add_data_requirement(self._requirement(wide)) + expected = np.concatenate([reader[index]["pair_margin"] for index in range(6)]) + self.assertTrue((expected < 1.0).all()) + batch = reader.decode_batch(list(range(6))) + self.assertEqual(float(batch["find_pair_margin"]), 1.0) + self.assertEqual(batch["pair_margin"].shape, (6, 1)) + np.testing.assert_allclose( + batch["pair_margin"].reshape(-1), expected, rtol=1e-12, atol=1e-12 + ) + + def test_a_threshold_keeps_the_side_of_every_frame(self) -> None: + """The window each frame is judged against is its own radius.""" + wide = 10.0 + distances = self._margins(wide) * wide + fraction = float(np.median(distances)) + margins = self._margins(fraction) + self.assertTrue((margins < 1.0).any() and (margins >= 1.0).any()) + np.testing.assert_array_equal(margins < 1.0, distances < fraction) + + def test_a_covalent_window_sizes_each_pair_by_its_elements(self) -> None: + """Oxygen and hydrogen carry different windows on the same geometry.""" + reader = LmdbDataReader(self._path, ["O", "H"], batch_size=6) + reader.add_data_requirement(self._requirement(0.5, "covalent")) + half = pair_half_thresholds(0.5, "covalent", ["O", "H"]) + for index in range(6): + frame = reader[index] + box = None if not frame["box"].any() else frame["box"] + np.testing.assert_allclose( + frame["pair_margin"], + [ + compute_min_pair_margin_single( + frame["coord"], box, frame["atype"], half + ) + ], + rtol=1e-12, + ) + # Half of the Pyykko single-bond radii puts the three pair thresholds + # at 0.63 Å for O-O, 0.475 Å for O-H and 0.32 Å for H-H, and an + # absolute window is one threshold for all three: no radius reproduces + # the covalent margins, which stay between the two extreme ones. + covalent = self._frame_margins(0.5, "covalent") + for threshold in (0.32, 0.475, 0.63): + self.assertFalse(np.allclose(covalent, self._frame_margins(threshold))) + self.assertTrue((covalent >= self._frame_margins(0.63)).all()) + self.assertTrue((covalent <= self._frame_margins(0.32)).all()) + + if __name__ == "__main__": unittest.main() diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py index 7739cdd9f4..9b61ee26cf 100644 --- a/source/tests/common/dpmodel/test_zbl_bridging.py +++ b/source/tests/common/dpmodel/test_zbl_bridging.py @@ -70,10 +70,12 @@ def test_builder_composes_linear_model(): kinds == ["EnergyAtomicModel", "InnerPotentialAtomicModel"] or kinds[1] == "InnerPotentialAtomicModel" ) - # radii wired to the LEARNED child's descriptor InnerClamp + # radii wired to the LEARNED child's descriptor: the switch spans the + # window and the clamp freezes two fifths into it 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 + assert float(dp_child.descriptor.bridging_switch.f_inner) == 0.8 + assert float(dp_child.descriptor.bridging_clamp.f_freeze) == pytest.approx(0.96) + assert dp_child.descriptor.bridging_scale == "absolute" def test_bridged_preset_belongs_to_output_statistics() -> None: @@ -294,12 +296,12 @@ def _pair_energy(model, natoms=2, r=1.0): class TestInnerPotentialChangeTypeMap: - """``change_type_map`` must rebuild the ZBL element lookup. + """``change_type_map`` must rebuild the ZBL coefficient tables. The generic ``BaseAtomicModel.change_type_map`` only rewrites the public - map and the stat/exclusion state; the nuclear-charge table belongs to - ``InnerPotential`` and is rebuilt there (review 3649295675). Without it - the lookup keeps the ORIGINAL elements while ``atype`` values already mean + map and the stat/exclusion state; the per-pair coefficients belong to + ``InnerPotential`` and are rebuilt there (review 3649295675). Without it + the tables keep the ORIGINAL elements while ``atype`` values already mean the new ones -- silently wrong energies, or ``IndexError`` for a longer map. """ @@ -311,17 +313,22 @@ def test_reorder_matches_a_freshly_built_model(self) -> None: fresh = InnerPotentialAtomicModel(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 + # table 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] + np.testing.assert_allclose( + np.asarray(model.potential.series_table), + np.asarray(fresh.potential.series_table), + rtol=1e-12, + ) assert model.potential.type_map == ["O", "H"] def test_added_element_extends_the_lookup(self) -> None: model = InnerPotentialAtomicModel(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] + assert model.potential.ntypes == 3 + # one row per ordered pair of the three types plus the padding type + assert model.potential.series_table.shape == (16, 8) # the new type is now addressable -- a stale (length-2) table raises # IndexError here graph_atype = np.full(2, 2, dtype=np.int64) @@ -344,8 +351,14 @@ def test_added_element_extends_the_lookup(self) -> None: def test_dropped_element_shrinks_the_lookup(self) -> None: model = InnerPotentialAtomicModel(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] + assert model.potential.ntypes == 1 + fresh = InnerPotentialAtomicModel(type_map=["Ni"], rcut=4.0, sel=[8]) + np.testing.assert_allclose( + np.asarray(model.potential.series_table), + np.asarray(fresh.potential.series_table), + rtol=1e-12, + ) + np.testing.assert_allclose(_pair_energy(model), _pair_energy(fresh), rtol=1e-12) def test_serialize_roundtrip_after_change_type_map(self) -> None: """Checkpoint continuity: the restored model must predict the same. @@ -770,8 +783,9 @@ def test_canonical_config_composes(self) -> None: # the composition derives the learned sibling's clamp window from # the inner_potential child: one source of truth for the radii 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 + assert float(dp_child.descriptor.bridging_switch.f_inner) == 0.8 + assert float(dp_child.descriptor.bridging_clamp.f_freeze) == pytest.approx(0.96) + assert dp_child.descriptor.bridging_scale == "absolute" assert dp_child.descriptor.bridging_switch is not None def test_canonical_matches_sugar_energy(self) -> None: @@ -963,9 +977,74 @@ def test_non_dpa4_learned_sibling_raises_cleanly(self) -> None: cfg = _canonical_config() cfg["models"][0]["descriptor"] = { "type": "se_e2_a", - "rcut": 4.0, "rcut_smth": 3.5, + "rcut": 4.0, "sel": [8, 8], } with pytest.raises(NotImplementedError, match="DPA4/SeZM"): get_model(cfg) + + +# --------------------------------------------------------------------------- +# The same composition in NLH mode +# --------------------------------------------------------------------------- +def _nlh_config(): + cfg = copy.deepcopy(ZBL_CONFIG) + cfg["bridging_method"] = "nlh" + return cfg + + +def test_nlh_bridging_builds_and_carries_the_mode(): + model = get_model(_nlh_config()) + assert isinstance(model, LinearEnergyModel) + potential = model.atomic_model.models[1].potential + assert potential.mode == "NLH" + + +def test_nlh_bridging_serialize_roundtrip_energy_identical(): + """The mode survives the wire and the restored model predicts the same.""" + model = get_model(_nlh_config()) + coord, atype, box = _close_pair_inputs() + data = model.serialize() + assert data["models"][1]["mode"] == "NLH" + restored = BaseModel.deserialize(data) + assert restored.atomic_model.models[1].potential.mode == "NLH" + e1 = model.call_common(coord, atype, box=box, neighbor_graph_method="dense")[ + "energy_redu" + ] + e2 = restored.call_common(coord, atype, box=box, neighbor_graph_method="dense")[ + "energy_redu" + ] + np.testing.assert_allclose(e1, e2, rtol=1e-12) + + +def test_nlh_and_zbl_bridging_predict_differently(): + """The mode reaches the energy; it is not silently dropped to the default.""" + coord, atype, box = _close_pair_inputs() + energies = [] + for cfg in (copy.deepcopy(ZBL_CONFIG), _nlh_config()): + model = get_model(cfg) + energies.append( + float( + model.call_common(coord, atype, box=box, neighbor_graph_method="dense")[ + "energy_redu" + ].sum() + ) + ) + assert abs(energies[0] - energies[1]) > 1e-6 + + +def test_nlh_atomic_model_roundtrips_the_mode(): + """The child atomic model carries the mode across serialize/deserialize.""" + child = InnerPotentialAtomicModel( + type_map=["Ni", "O"], mode="nlh", rcut=4.0, sel=[8] + ) + data = child.serialize() + assert data["mode"] == "NLH" + restored = InnerPotentialAtomicModel.deserialize(data) + assert restored.potential.mode == "NLH" + np.testing.assert_allclose( + np.asarray(restored.potential.series_table), + np.asarray(child.potential.series_table), + rtol=1e-12, + ) diff --git a/source/tests/common/test_bridging.py b/source/tests/common/test_bridging.py index 24480ffddf..6432d46955 100644 --- a/source/tests/common/test_bridging.py +++ b/source/tests/common/test_bridging.py @@ -8,11 +8,25 @@ """ import copy +import json +from pathlib import ( + Path, +) +import numpy as np import pytest +from deepmd.dpmodel.descriptor.dpa4_nn import ( + BridgingSwitch, +) from deepmd.utils.bridging import ( + DEFAULT_FRACTION_INNER, + DEFAULT_FRACTION_OUTER, expand_bridging_method, + resolve_bridging_window, +) +from deepmd.utils.data import ( + min_pair_dist_requirement, ) @@ -58,22 +72,51 @@ def test_expansion_shape() -> None: assert inner == { "type": "inner_potential", "mode": "ZBL", + "fraction_inner": 0.26, + "fraction_outer": 0.80, "r_inner": 0.8, "r_outer": 1.2, } # the flag keys must not leak into the canonical config - for key in ("bridging_method", "bridging_r_inner", "bridging_r_outer"): + for key in ( + "bridging_method", + "bridging_fraction_inner", + "bridging_fraction_outer", + "bridging_r_inner", + "bridging_r_outer", + ): assert key not in out assert key not in learned -def test_default_radii() -> None: +def test_window_defaults_to_the_covalent_rule() -> None: + """Without explicit radii the window is a fraction of each pair's own size.""" data = _flag_config() del data["bridging_r_inner"] del data["bridging_r_outer"] inner = expand_bridging_method(data)["models"][1] - assert inner["r_inner"] == 0.5 - assert inner["r_outer"] == 0.8 + assert inner["r_inner"] is None + assert inner["r_outer"] is None + assert resolve_bridging_window(inner) == { + "inner_clamp_f_inner": 0.26, + "inner_clamp_f_outer": 0.80, + "inner_clamp_scale": "covalent", + } + + +def test_explicit_radii_override_the_covalent_rule() -> None: + """Radii in Å apply the same window to every pair, through a unit scale.""" + inner = expand_bridging_method(_flag_config())["models"][1] + assert resolve_bridging_window(inner) == { + "inner_clamp_f_inner": 0.8, + "inner_clamp_f_outer": 1.2, + "inner_clamp_scale": "absolute", + } + + +def test_half_a_window_is_rejected() -> None: + with pytest.raises(ValueError, match="both `r_inner` and `r_outer`"): + resolve_bridging_window({"r_outer": 1.2}) def test_input_is_not_mutated() -> None: @@ -345,3 +388,79 @@ def test_routing_resolves_argcheck_default_conflicts() -> None: route_canonical_learned_options( {"data_stat_protect": 0.2}, {"data_stat_protect": 0.3} ) + + +# The frame filter of a bridged model is the midpoint of its window, so the two +# decisions it rests on -- which models count as bridged, and which radius the +# window yields -- are pinned here against the shipped examples. + +_EXAMPLES = Path(__file__).parent.parent.parent.parent / "examples" / "water" +_BRIDGED_INPUTS = ( + _EXAMPLES / "dpa4" / "input-zbl.json", + _EXAMPLES / "dpa4c" / "input-zbl.json", +) + + +@pytest.mark.parametrize("path", _BRIDGED_INPUTS, ids=lambda p: p.parent.name) +def test_a_shipped_bridged_example_derives_its_own_filter(path) -> None: + """Every learned family the composition accepts owns a frame filter.""" + config = json.loads(path.read_text()) + training_data = config["training"]["training_data"] + assert "min_pair_dist" not in training_data, ( + "a bridged model derives the filter from its window, so the input must " + "not carry a second copy of it" + ) + requirement = min_pair_dist_requirement(config["model"], 0.0) + assert requirement is not None + assert requirement.key == "pair_margin" + assert requirement.default == pytest.approx( + 0.5 * (DEFAULT_FRACTION_INNER + DEFAULT_FRACTION_OUTER) + ) + assert requirement.length_scale == "covalent" + + +@pytest.mark.parametrize("path", _BRIDGED_INPUTS, ids=lambda p: p.parent.name) +def test_a_shipped_bridged_example_refuses_an_explicit_filter(path) -> None: + config = json.loads(path.read_text()) + with pytest.raises(ValueError, match="must not be set on a bridged model"): + min_pair_dist_requirement(config["model"], 0.5) + + +def test_an_explicit_window_filters_at_its_midpoint() -> None: + """Radii given in Å yield a filter in Å, halfway between the two.""" + requirement = min_pair_dist_requirement(_flag_config(), 0.0) + assert requirement is not None + assert requirement.default == pytest.approx(1.0) + assert requirement.length_scale == "absolute" + + +@pytest.mark.parametrize( + "fractions", + [(DEFAULT_FRACTION_INNER, DEFAULT_FRACTION_OUTER), (0.1, 0.9), (0.5, 0.8)], +) +def test_the_bridging_switch_is_half_open_at_the_filter_radius(fractions) -> None: + """The filter keeps exactly the pairs whose learned amplitude is at least one half.""" + f_inner, f_outer = fractions + config = _flag_config() + del config["bridging_r_inner"], config["bridging_r_outer"] + config["bridging_fraction_inner"] = f_inner + config["bridging_fraction_outer"] = f_outer + requirement = min_pair_dist_requirement(config, 0.0) + assert requirement is not None + contact = np.array([1.7]) + opening = BridgingSwitch(f_inner, f_outer).call( + requirement.default * contact, contact + ) + np.testing.assert_allclose(opening, 0.5, rtol=0.0, atol=1.0e-14) + + +def test_an_unbridged_model_without_a_floor_registers_no_filter() -> None: + """Whether the filter runs and whether it is registered is one decision.""" + config = _flag_config() + for key in ("bridging_method", "bridging_r_inner", "bridging_r_outer"): + del config[key] + assert min_pair_dist_requirement(config, 0.0) is None + requirement = min_pair_dist_requirement(config, 0.5) + assert requirement is not None + assert requirement.default == 0.5 + assert requirement.length_scale == "absolute" diff --git a/source/tests/common/test_deepmd_data.py b/source/tests/common/test_deepmd_data.py index 0f26d2466a..9e9e57d9eb 100644 --- a/source/tests/common/test_deepmd_data.py +++ b/source/tests/common/test_deepmd_data.py @@ -1,4 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import json +import subprocess +import sys import tempfile import unittest from pathlib import ( @@ -7,9 +10,16 @@ import numpy as np +from deepmd.dpmodel.utils.dist_check import ( + compute_min_pair_margin_single, + pair_half_thresholds, +) from deepmd.utils.data import ( DeepmdData, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) class TestDeepmdDataTypeMap(unittest.TestCase): @@ -47,6 +57,70 @@ def test_remap_with_unused_types(self) -> None: expected_sorted = expected_atom_types[data.idx_map] np.testing.assert_array_equal(loaded["type"], np.tile(expected_sorted, (1, 1))) + def test_covalent_margin_without_dataset_type_map(self) -> None: + """Model type indices retain their element names without a dataset map.""" + (self.root / "type_map.raw").unlink() + coord = np.array( + [[0.0, 0.0, 0.0], [0.4, 0.0, 0.0], [2.0, 0.0, 0.0], [4.0, 0.0, 0.0]] + ) + box = 6.0 * np.eye(3) + np.save(self.set_dir / "coord.npy", coord.reshape(1, -1)) + np.save(self.set_dir / "box.npy", box.reshape(1, 9)) + type_map = ["O", "H", "Si"] + data = DeepmdData(str(self.root), type_map=type_map) + data.add("pair_margin", 1, default=0.5, length_scale="covalent") + expected = compute_min_pair_margin_single( + coord, + box, + np.array([0, 1, 0, 1]), + pair_half_thresholds(0.5, "covalent", type_map), + screened=True, + ) + for loaded in (data.get_single_frame(0, num_worker=1), data.get_batch(1)): + self.assertEqual(float(loaded["find_pair_margin"]), 1.0) + np.testing.assert_allclose(loaded["pair_margin"], expected) + self.assertIsNone(data.get_type_map()) + self.assertEqual(data.get_ntypes(), 2) + np.testing.assert_array_equal(data.atom_type, [0, 1, 0, 1]) + + def test_required_dataset_type_map_cannot_be_replaced(self) -> None: + """A supplied model map does not make a required dataset map optional.""" + (self.root / "type_map.raw").unlink() + with self.assertRaisesRegex(AssertionError, "must have type_map.raw"): + DeepmdData(str(self.root), type_map=["O", "H"], optional_type_map=False) + + def test_model_map_preserves_virtual_spin_type_counts(self) -> None: + """Unnamed virtual types remain part of the reader's numeric type domain.""" + fixture_dir = Path(__file__).parent.parent / "tf" + model = json.loads((fixture_dir / "test_model_spin.json").read_text())["model"] + system = fixture_dir / "model_spin" + self.assertFalse((system / "type_map.raw").exists()) + stored_types = np.loadtxt(system / "type.raw", dtype=np.int32) + counts = np.bincount(stored_types) + self.assertGreater(len(counts), len(model["type_map"])) + data = DeepmdDataSystem( + [str(system)], batch_size=1, test_size=1, type_map=model["type_map"] + ) + self.assertEqual(data.get_ntypes(), len(counts)) + np.testing.assert_array_equal( + data.get_batch()["natoms_vec"], + np.concatenate(([stored_types.size, stored_types.size], counts)), + ) + + def test_data_system_imports_without_preloading_dpmodel(self) -> None: + """The generic reader imports independently of backend initialization.""" + result = subprocess.run( + [ + sys.executable, + "-c", + "from deepmd.utils.data_system import DeepmdDataSystem", + ], + capture_output=True, + text=True, + check=False, + ) + self.assertEqual(result.returncode, 0, result.stderr) + class _RaisingModifier: """A modifier whose ``modify_data`` always fails.""" diff --git a/source/tests/common/test_dpmodel_train.py b/source/tests/common/test_dpmodel_train.py index 22149f8e99..0951090293 100644 --- a/source/tests/common/test_dpmodel_train.py +++ b/source/tests/common/test_dpmodel_train.py @@ -1,5 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import unittest +from collections.abc import ( + Callable, +) +from typing import ( + Any, +) import numpy as np @@ -33,8 +39,8 @@ def change_out_bias(self, sample_func, bias_adjust_mode): def get_type_map(self): return ["O", "H"] - def get_fitting_net(self): - return self.fitting_net + def compute_fitting_input_stats(self, sample_func: Callable[[], Any]) -> None: + self.fitting_net.compute_input_stats(sample_func) class TestChangeModelOutBias(unittest.TestCase): diff --git a/source/tests/common/test_examples.py b/source/tests/common/test_examples.py index 416ddb2b60..ff11aff68d 100644 --- a/source/tests/common/test_examples.py +++ b/source/tests/common/test_examples.py @@ -72,6 +72,7 @@ p_examples / "water" / "dpa4" / "input-zbl.json", p_examples / "water" / "dpa4" / "lmp" / "input.json", p_examples / "water" / "dpa4c" / "input.json", + p_examples / "water" / "dpa4c" / "input-zbl.json", p_examples / "property" / "train" / "input_torch.json", p_examples / "water" / "se_e3_tebd" / "input_torch.json", p_examples / "hessian" / "single_task" / "input.json", diff --git a/source/tests/common/test_spin.py b/source/tests/common/test_spin.py index fa8c91a530..84acc2684e 100644 --- a/source/tests/common/test_spin.py +++ b/source/tests/common/test_spin.py @@ -6,12 +6,45 @@ from deepmd.utils.spin import ( Spin, + env_protection_descriptors, normalize_spin_use_spin, ) CUR_DIR = os.path.dirname(__file__) +class EnvProtectionDescriptorsTest(unittest.TestCase): + def test_hybrid_only_configures_environment_matrix_children(self) -> None: + environment = {"type": "se_e2_a", "env_protection": 0.02} + dpa4 = {"type": "DPA4", "eps": 1e-7} + nested_environment = {"type": "se_atten"} + descriptor = { + "type": "hybrid", + "env_protection": 0.03, + "list": [ + environment, + dpa4, + { + "type": "hybrid", + "env_protection": 0.04, + "list": [nested_environment, {"type": "dpa4c"}], + }, + ], + } + selected = list(env_protection_descriptors(descriptor)) + self.assertEqual(len(selected), 2) + self.assertIs(selected[0], environment) + self.assertIs(selected[1], nested_environment) + for child in selected: + child.setdefault("env_protection", 0.01) + self.assertEqual(environment["env_protection"], 0.02) + self.assertEqual(nested_environment["env_protection"], 0.04) + self.assertNotIn("env_protection", descriptor) + self.assertNotIn("env_protection", descriptor["list"][2]) + self.assertEqual(dpa4, {"type": "DPA4", "eps": 1e-7}) + self.assertEqual(list(env_protection_descriptors({"type": "SeZM"})), []) + + class NormalizeUseSpinTest(unittest.TestCase): """Unit tests for ``normalize_spin_use_spin`` (pure; all three forms).""" diff --git a/source/tests/infer/gen_dpa4_spin_zbl.py b/source/tests/infer/gen_dpa4_spin_zbl.py index eb191f0104..13f236dd49 100644 --- a/source/tests/infer/gen_dpa4_spin_zbl.py +++ b/source/tests/infer/gen_dpa4_spin_zbl.py @@ -23,9 +23,11 @@ Multi-rank capable, and this fixture PINS that contract ------------------------------------------------------- -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 the per-node partials are +Bridging enables the descriptor's Source Freeze Propagation Gate, which folds +into the envelope of an edge the amplitudes ``w`` of every pair of that edge's +source OTHER than its own, and hands the full per-node product +``eta_j = prod_{e: src_e = j} w_e`` to the model as the readout gate. +Edges exist only for owned centres, so the per-node partials are rank-incomplete; the with-comm artifact completes them with one reverse-accumulate + forward-broadcast border exchange before the gate is applied (issue #5906). ``_check_metadata`` asserts BOTH diff --git a/source/tests/infer/gen_dpa4_zbl.py b/source/tests/infer/gen_dpa4_zbl.py index c4bbbde6f9..a711eb81eb 100644 --- a/source/tests/infer/gen_dpa4_zbl.py +++ b/source/tests/infer/gen_dpa4_zbl.py @@ -11,8 +11,10 @@ never touches ``DeepPotPTExpt``. Multi-rank capable (issue #5906): 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, +Freeze Propagation Gate, which folds into the envelope of an edge the +amplitudes ``w`` of every pair of that edge's source OTHER than its own, and +hands the full per-node product ``eta_j = prod_{e: src_e = j} w_e`` to the +model as the readout gate. Edges exist only for owned centres, so the per-node ``[log_eta, zero_count]`` partials are rank-incomplete; the with-comm artifact completes them via one reverse-accumulate (``border_op_backward``) + forward-broadcast (``border_op``) exchange before diff --git a/source/tests/pt/model/test_descriptor_sezm.py b/source/tests/pt/model/test_descriptor_sezm.py index b3a43a0ae4..5bf465da0f 100644 --- a/source/tests/pt/model/test_descriptor_sezm.py +++ b/source/tests/pt/model/test_descriptor_sezm.py @@ -596,7 +596,7 @@ def test_zero_block_descriptor(self) -> None: ) edge_vec = flat[edge_index[0]] - flat[edge_index[1]] edge_mask = torch.ones(2, dtype=torch.bool, device=self.device) - desc_e, latent, _ = model.forward_with_edges( + desc_e, latent, _, _ = model.forward_with_edges( extended_coord=coord.reshape(1, -1), extended_atype=atype, edge_index=edge_index, @@ -1074,7 +1074,7 @@ def test_charge_spin_sparse_edge_conditioning(self) -> None: nlist, charge_spin=torch.tensor([[0.0, 1.0]], device=self.device), ) - desc_ref, _, _ = model.forward_with_edges( + desc_ref, _, _, _ = model.forward_with_edges( extended_coord=coord, extended_atype=atype, edge_index=edge_index, @@ -1082,7 +1082,7 @@ def test_charge_spin_sparse_edge_conditioning(self) -> None: edge_mask=edge_mask, charge_spin=torch.tensor([[0.0, 1.0]], device=self.device), ) - desc_shifted, _, _ = model.forward_with_edges( + desc_shifted, _, _, _ = model.forward_with_edges( extended_coord=coord, extended_atype=atype, edge_index=edge_index, @@ -1091,7 +1091,7 @@ def test_charge_spin_sparse_edge_conditioning(self) -> None: charge_spin=torch.tensor([[1.0, 1.0]], device=self.device), ) restored = DescrptSeZM.deserialize(model.serialize()) - desc_restored, _, _ = restored.forward_with_edges( + desc_restored, _, _, _ = restored.forward_with_edges( extended_coord=coord, extended_atype=atype, edge_index=edge_index, @@ -1291,7 +1291,7 @@ def test_descriptor_spin_joint_rotation_invariance(self) -> None: p.copy_(torch.randn_like(p) * 0.1) model.eval() - desc, _, _ = model.forward_with_edges( + desc, _, _, _ = model.forward_with_edges( extended_coord=coord, extended_atype=atype, edge_index=edge_index, @@ -1299,7 +1299,7 @@ def test_descriptor_spin_joint_rotation_invariance(self) -> None: edge_mask=edge_mask, spin=spin, ) - desc_rot, _, _ = model.forward_with_edges( + desc_rot, _, _, _ = model.forward_with_edges( extended_coord=coord, extended_atype=atype, edge_index=edge_index, @@ -1310,7 +1310,7 @@ def test_descriptor_spin_joint_rotation_invariance(self) -> None: torch.testing.assert_close(desc, desc_rot, atol=1e-9, rtol=1e-9) # Spin actually changes the descriptor (injection is not a no-op). - desc_zero, _, _ = model.forward_with_edges( + desc_zero, _, _, _ = model.forward_with_edges( extended_coord=coord, extended_atype=atype, edge_index=edge_index, @@ -1383,7 +1383,7 @@ def test_zero_gate_receives_a_gradient(self) -> None: """ model = self._descriptor() kwargs, spin = self._inputs() - desc, _, _ = model.forward_with_edges(**kwargs, spin=spin) + desc, _, _, _ = model.forward_with_edges(**kwargs, spin=spin) desc.sum().backward() gate_grad = model.env_seed_embedding.spin_scale.grad self.assertIsNotNone(gate_grad) @@ -1402,10 +1402,10 @@ def test_gate_placement_differs_from_the_legacy_one_by_a_square(self) -> None: amplitude = 2.0 with torch.no_grad(): model.env_seed_embedding.spin_scale.fill_(amplitude**2) - migrated, _, _ = model.forward_with_edges(**kwargs, spin=spin) + migrated, _, _, _ = model.forward_with_edges(**kwargs, spin=spin) with torch.no_grad(): model.env_seed_embedding.spin_scale.fill_(1.0) - legacy, _, _ = model.forward_with_edges(**kwargs, spin=amplitude * spin) + legacy, _, _, _ = model.forward_with_edges(**kwargs, spin=amplitude * spin) torch.testing.assert_close(migrated, legacy, atol=1e-12, rtol=1e-12) def test_loading_a_legacy_state_squares_the_gate(self) -> None: @@ -1427,8 +1427,8 @@ def test_loading_a_legacy_state_squares_the_gate(self) -> None: model.env_seed_embedding.spin_scale.detach(), torch.full_like(model.env_seed_embedding.spin_scale, 4.0), ) - self.assertEqual(model.version, 1.2) - self.assertEqual(float(model.version_tensor.item()), 1.2) + self.assertEqual(model.version, 1.3) + self.assertEqual(float(model.version_tensor.item()), 1.3) def test_loading_legacy_spin_free_state_zeros_dormant_routes(self) -> None: kwargs = {**self._kwargs(), "use_spin": [False, False]} @@ -1456,10 +1456,10 @@ def test_loading_legacy_spin_free_state_zeros_dormant_routes(self) -> None: migrated["spin_embedding.mag_layer1.matrix"], torch.full_like(migrated["spin_embedding.mag_layer1.matrix"], 5.0), ) - self.assertEqual(model.version, 1.2) + self.assertEqual(model.version, 1.3) def test_a_migrated_state_is_migrated_only_once(self) -> None: - """Re-saving a migrated descriptor advertises 1.2, so a reload is inert.""" + """A migrated descriptor advertises its current format; a reload is inert.""" state = self._descriptor().state_dict() state["version_tensor"] = torch.full_like(state["version_tensor"], 1.1) state["env_seed_embedding.spin_scale"] = torch.full_like( @@ -1473,7 +1473,7 @@ def test_a_migrated_state_is_migrated_only_once(self) -> None: reloaded.env_seed_embedding.spin_scale.detach(), torch.full_like(reloaded.env_seed_embedding.spin_scale, 4.0), ) - self.assertEqual(reloaded.version, 1.2) + self.assertEqual(reloaded.version, 1.3) def test_serialize_roundtrip_migrates_a_legacy_payload(self) -> None: """The dp-format path shares the rule with the state-dict path.""" @@ -1487,7 +1487,7 @@ def test_serialize_roundtrip_migrates_a_legacy_payload(self) -> None: model.env_seed_embedding.spin_scale.detach(), torch.full_like(model.env_seed_embedding.spin_scale, 9.0), ) - self.assertEqual(model.version, 1.2) + self.assertEqual(model.version, 1.3) class TestBuildEdgeQuaternion(_SeZMTestCase): @@ -2105,114 +2105,117 @@ def test_node_engine_rotation_equivariance(self) -> None: class TestInnerClamp(_SeZMTestCase): - """Test InnerClamp C3-continuous septic Hermite clamping.""" + """The distance clamp of a bridging window. + + Below the freeze point the distance the descriptor is shown stays at that + point, above the outer radius the clamp is the identity, and the septic + Hermite profile joins the two with three continuous derivatives. The + freeze point sits two fifths into the window, just below the midpoint + where the frame filter ends the training data, so the shown distance + still moves at the filter point. + """ def setUp(self) -> None: super().setUp() - self.r_inner = 1.0 + # A unit pair length scale makes the two fractions the window radii in + # Å directly, so every boundary below is read at the plain distance. + self.r_inner = 0.5 self.r_outer = 1.5 + self.r_freeze = 0.9 + self.contact = torch.tensor(1.0, dtype=torch.float64, device=self.device) self.clamp = InnerClamp(self.r_inner, self.r_outer) + self.assertAlmostEqual(self.clamp.f_freeze, self.r_freeze, places=14) + + def _derivatives(self, r: torch.Tensor, order: int) -> list[torch.Tensor]: + """Successive derivatives of the clamp, from the first to ``order``.""" + r = r.detach().clone().requires_grad_(True) + value = self.clamp(r, self.contact) + derivatives = [] + for step in range(order): + value = torch.autograd.grad(value.sum(), r, create_graph=step + 1 < order)[ + 0 + ] + derivatives.append(value.detach()) + return derivatives def test_monotonicity(self) -> None: - """Test that r̃ is monotonically non-decreasing.""" + """The shown distance never decreases as the true distance grows.""" r = torch.linspace(0.0, 3.0, 1000, dtype=torch.float64, device=self.device) - out = self.clamp(r) + out = self.clamp(r, self.contact) diff = out[1:] - out[:-1] self.assertTrue((diff >= -1e-14).all(), "InnerClamp is not monotonic") - def test_frozen_zone_zero_gradient(self) -> None: - """Test that dr̃/dr = 0 for r < r_inner (frozen zone).""" + def test_frozen_zone_sits_at_the_freeze_point(self) -> None: + """Below the freeze point the shown distance is that point itself.""" r = torch.tensor( - [0.3, 0.5, 0.8, 0.99], + [0.0, 0.3, 0.5, 0.8, 0.89], dtype=torch.float64, device=self.device, - requires_grad=True, ) - out = self.clamp(r) - grads = torch.autograd.grad(out.sum(), r)[0] + torch.testing.assert_close( + self.clamp(r, self.contact), + torch.full_like(r, self.r_freeze), + atol=1e-14, + rtol=0, + ) + grads = self._derivatives(r, 1)[0] torch.testing.assert_close( grads, torch.zeros_like(grads), - atol=1e-12, + atol=1e-14, rtol=0, - msg="Gradient should be zero in the frozen zone", + msg="the frozen zone must carry no gradient", ) def test_identity_zone_unit_gradient(self) -> None: - """Test that dr̃/dr = 1 for r > r_outer (identity zone).""" + """Beyond the outer radius the clamp returns the distance unchanged.""" r = torch.tensor( - [1.6, 2.0, 3.0, 5.0], + [1.5, 1.6, 2.0, 3.0, 5.0], dtype=torch.float64, device=self.device, - requires_grad=True, ) - out = self.clamp(r) - grads = torch.autograd.grad(out.sum(), r)[0] + torch.testing.assert_close(self.clamp(r, self.contact), r, atol=0, rtol=0) + grads = self._derivatives(r, 1)[0] torch.testing.assert_close( grads, torch.ones_like(grads), - atol=1e-12, + atol=1e-14, rtol=0, - msg="Gradient should be 1 in the identity zone", + msg="the identity zone must carry a unit gradient", ) def test_c3_continuity_at_boundaries(self) -> None: - """Test C3 continuity at r_inner and r_outer via autograd derivatives.""" - eps = 1e-6 - for boundary in [self.r_inner, self.r_outer]: + """Three derivatives match across both ends of the transition. + + The septic profile vanishes together with its first two derivatives at + the freeze point and reaches unit slope with vanishing curvature at the + outer radius, so the clamp joins the frozen and the identity branch + with three continuous derivatives. The step ``eps`` is small enough + that the genuine variation of the third derivative over it stays far + below the tolerance, which leaves a real discontinuity as the only way + to fail. + """ + eps = 1.0e-9 + for boundary in (self.r_freeze, self.r_outer): r = torch.tensor( [boundary - eps, boundary, boundary + eps], dtype=torch.float64, device=self.device, - requires_grad=True, - ) - out = self.clamp(r) - - # First derivative via autograd - grads = torch.autograd.grad(out.sum(), r, create_graph=True)[0] - # dr̃/dr should be continuous (left ≈ center ≈ right) - self.assertAlmostEqual( - grads[0].item(), - grads[1].item(), - places=4, - msg=f"First derivative discontinuous at {boundary}", - ) - self.assertAlmostEqual( - grads[1].item(), - grads[2].item(), - places=4, - msg=f"First derivative discontinuous at {boundary}", - ) - - # Second derivative via autograd - grads2 = torch.autograd.grad(grads.sum(), r, create_graph=True)[0] - self.assertAlmostEqual( - grads2[0].item(), - grads2[1].item(), - places=3, - msg=f"Second derivative discontinuous at {boundary}", - ) - self.assertAlmostEqual( - grads2[1].item(), - grads2[2].item(), - places=3, - msg=f"Second derivative discontinuous at {boundary}", - ) - - # Third derivative via autograd - grads3 = torch.autograd.grad(grads2.sum(), r)[0] - self.assertAlmostEqual( - grads3[0].item(), - grads3[1].item(), - places=2, - msg=f"Third derivative discontinuous at {boundary}", - ) - self.assertAlmostEqual( - grads3[1].item(), - grads3[2].item(), - places=2, - msg=f"Third derivative discontinuous at {boundary}", ) + for order, values in enumerate(self._derivatives(r, 3), start=1): + left, center, right = (float(v) for v in values) + self.assertAlmostEqual( + left, + center, + delta=1e-4, + msg=f"derivative {order} jumps below {boundary}", + ) + self.assertAlmostEqual( + center, + right, + delta=1e-4, + msg=f"derivative {order} jumps above {boundary}", + ) def test_invalid_params(self) -> None: """Test that invalid parameters raise ValueError.""" @@ -2311,49 +2314,6 @@ def test_envelope_nextafter_cutoff_attention(self) -> None: ) self.assertLess(float(alpha[1, 0, 0]), 1.0e-30) - def test_tiny_source_weight_hessian(self) -> None: - """Log-domain source scaling must preserve the physical Hessian.""" - logits = torch.tensor( - [[[0.0]], [[20.0]]], dtype=torch.float32, device=self.device - ) - edge_env = torch.ones((2, 1), dtype=torch.float32, device=self.device) - dst = torch.zeros(2, dtype=torch.int64, device=self.device) - z_bias_raw = torch.tensor( - [[math.log(math.expm1(1.0))]], dtype=torch.float32, device=self.device - ) - eps = 1.0e-7 - - def attention_sum(source_weight: torch.Tensor) -> torch.Tensor: - return segment_envelope_gated_softmax( - logits, - edge_env, - dst, - 1, - z_bias_raw, - eps, - source_weight[:, None], - ).sum() - - null_mass = torch.nn.functional.softplus(z_bias_raw[0, 0]) + eps - - def physical_sum(source_weight: torch.Tensor) -> torch.Tensor: - edge_mass = source_weight * torch.exp(logits[:, 0, 0]) - return (edge_mass / (null_mass + edge_mass.sum())).sum() - - source_weight = torch.tensor( - [1.0, 1.0e-30], dtype=torch.float32, device=self.device - ) - actual = torch.autograd.functional.hessian(attention_sum, source_weight) - reference = torch.autograd.functional.hessian(physical_sum, source_weight) - self.assertTrue(bool(torch.isfinite(actual).all())) - torch.testing.assert_close( - actual[0, 0], - reference[0, 0], - rtol=1.0e-5, - atol=1.0e-6, - ) - torch.testing.assert_close(actual, reference, rtol=1.0e-5, atol=32.0) - class TestEdgeNorm(_SeZMTestCase): """The ``edge_norm`` switch and its effect on cutoff smoothness. @@ -2613,7 +2573,6 @@ def _build_model_params( "channels": 16, "n_focus": n_focus, "focus_dim": 0, - "focus_compete": True, "n_radial": 6, "radial_mlp": [16], "use_env_seed": True, @@ -2974,10 +2933,15 @@ class TestSourceFreezePropagationGate(TestDescriptorEnergyCurveSmoothness): ``r_hat_BC`` — these are exactly the channels through which the direction and multi-hop leaks manifest in the user's analysis. - The guarantee SFPG delivers is therefore very sharp: **the atomic - energy of every node that is not ``B`` itself must stay strictly - constant**. ``E_B`` is expected to vary because ``B`` still has a - genuine chemical bond with ``C``. + The gate leaves the two edges between ``A`` and ``B`` alive, so the + descriptors of the frozen pair still follow each other at the clamped + distance; what leaves the graph are the edges the pair emits into + ``C``, and what closes the remaining channel is the per-node gate, + which is zero on ``A`` and on ``B``. The guarantee is therefore sharp + at the level of the energy: **the atomic energy of ``A`` is the per-type + bias of the fitting plus the analytical halves of its own pairs, and + stays strictly constant** along the trajectory. ``E_B`` moves only + through the analytical term of its genuine bond with ``C``. """ NEAR_DISTANCE = 0.6 # < BRIDGING_R_INNER, stays frozen for all samples diff --git a/source/tests/pt/model/test_descriptor_sezm_train_paths.py b/source/tests/pt/model/test_descriptor_sezm_train_paths.py index cdf53b1d5e..2629d64305 100644 --- a/source/tests/pt/model/test_descriptor_sezm_train_paths.py +++ b/source/tests/pt/model/test_descriptor_sezm_train_paths.py @@ -104,8 +104,11 @@ def _make_descriptor( random_gamma=False, precision=precision, seed=7, - inner_clamp_r_inner=0.8 if source_gated else None, - inner_clamp_r_outer=1.2 if source_gated else None, + # The absolute scale gives every pair the unit length scale, so the two + # fractions are the window radii in Å. + inner_clamp_f_inner=0.8 if source_gated else None, + inner_clamp_f_outer=1.2 if source_gated else None, + inner_clamp_scale="absolute", ) @@ -303,8 +306,15 @@ def _inference_step(self, descriptor: DescrptSeZM) -> tuple[np.ndarray, np.ndarr gradient = torch.autograd.grad(output.sum(), coord)[0] return output.detach().cpu().numpy(), gradient.detach().cpu().numpy() - def test_source_gated_flash_retains_dense_rotations(self, monkeypatch) -> None: - """Source-gated flash inference retains the rotations its fallback uses.""" + def test_source_gated_inference_takes_the_fused_convolution( + self, monkeypatch + ) -> None: + """A bridged descriptor runs the fused convolution and matches the dense path. + + The source gate is folded into the edge envelope, so the fused + convolution and the flash attention see it through the one factor they + already weight by, and the dense per-edge rotations are not built. + """ if not cuda_infer_available(): pytest.skip("the DPA4 CUDA inference operators are unavailable") if not SO2_VALUE_PATH_TRITON_AVAILABLE: @@ -332,8 +342,8 @@ def test_source_gated_flash_retains_dense_rotations(self, monkeypatch) -> None: pytest.skip("the descriptor layout has no fused CUDA convolution") assert conv._flash_atten_fn is not None assert conv._cuda_value_train is None - assert not accelerated._wigner_free_conv - assert accelerated._build_full_wigner() + assert accelerated._wigner_free_conv + assert not accelerated._build_full_wigner() output, gradient = self._inference_step(accelerated) np.testing.assert_allclose(output, dense_output, rtol=2e-4, atol=2e-5) diff --git a/source/tests/pt/model/test_dpa4_dpmodel_parity.py b/source/tests/pt/model/test_dpa4_dpmodel_parity.py index 898e1536fc..0c4e3ca7f2 100644 --- a/source/tests/pt/model/test_dpa4_dpmodel_parity.py +++ b/source/tests/pt/model/test_dpa4_dpmodel_parity.py @@ -2038,6 +2038,8 @@ def _build_so2_edge_data( a parity failure. ``edge_env`` is the exception: the production ``build_edge_cache`` multiplies the envelope by the slot mask, so it is exactly zero on invalid slots, and this fixture mirrors that contract. + ``with_gate`` multiplies a random source gate into the envelope, as the + production builder does for a bridged model. ``n_radial``: when not None, ``edge_rbf`` is filled with random values of width ``n_radial`` (garbage in masked slots too); otherwise it stays the @@ -2094,9 +2096,11 @@ def _build_so2_edge_data( # baseline aggregation (n_atten_head=0) relies on this; the attention path # masks independently, so both stay parity-correct. edge_env = rng.uniform(0.2, 1.0, size=(n_edge, 1)) * mask[:, None] + if with_gate: + # A bridged cache folds the source gate into the envelope. + edge_env = edge_env * rng.uniform(0.1, 1.0, size=(n_edge, 1)) deg = ((edge_env[:, 0] ** 2) * mask).reshape(nloc, nnei).sum(axis=1) inv_sqrt_deg = (1.0 / np.sqrt(deg + 1.0)).reshape(nloc, 1, 1) - edge_src_gate = rng.uniform(0.1, 1.0, size=(n_edge, 1)) if with_gate else None radial = rng.normal(size=(n_edge, lmax + 1, channels)) x = rng.normal(size=(nloc, dim_full, channels)) @@ -2112,7 +2116,6 @@ def _build_so2_edge_data( inv_sqrt_deg=t(inv_sqrt_deg), D_full=t(D_full[valid]), Dt_full=t(Dt_full[valid]), - edge_src_gate=None if edge_src_gate is None else t(edge_src_gate[valid]), ) dp_cache = EdgeCache( src=src, @@ -2125,7 +2128,6 @@ def _build_so2_edge_data( inv_sqrt_deg=inv_sqrt_deg, D_full=D_full, Dt_full=Dt_full, - edge_src_gate=edge_src_gate, edge_mask=mask, ) return pt_cache, dp_cache, radial, radial[valid], x, valid @@ -2318,8 +2320,8 @@ def test_radial_degree_mixer_errors(self) -> None: @pytest.mark.parametrize( "masked", ["none", "slots", "node"] ) # padded-slot patterns (node = one all-masked destination) - @pytest.mark.parametrize("use_src_weight", [False, True]) # SFPG gate branch - def test_segment_envelope_gated_softmax(self, masked, use_src_weight) -> None: + @pytest.mark.parametrize("with_gate", [False, True]) # folded SFPG gate + def test_segment_envelope_gated_softmax(self, masked, with_gate) -> None: from deepmd.dpmodel.descriptor.dpa4_nn.attention import ( segment_envelope_gated_softmax as dp_softmax, ) @@ -2336,7 +2338,7 @@ def test_segment_envelope_gated_softmax(self, masked, use_src_weight) -> None: lmax=2, channels=4, masked=masked, - with_gate=use_src_weight, + with_gate=with_gate, ) n_edge = nloc * nnei logits = rng.normal(size=(n_edge, n_focus, n_head)) @@ -2349,7 +2351,6 @@ def test_segment_envelope_gated_softmax(self, masked, use_src_weight) -> None: n_nodes=nloc, z_bias_raw=z_bias_raw, eps=1e-7, - src_weight=dp_cache.edge_src_gate, edge_mask=dp_cache.edge_mask, ) alpha_pt = pt_softmax( @@ -2359,7 +2360,6 @@ def test_segment_envelope_gated_softmax(self, masked, use_src_weight) -> None: n_nodes=nloc, z_bias_raw=to_pt(z_bias_raw), eps=1e-7, - src_weight=pt_cache.edge_src_gate, ) alpha_dp = np.asarray(alpha_dp) np.testing.assert_allclose( @@ -2419,6 +2419,48 @@ def evaluate(crossing_envelope: float) -> np.ndarray: ) assert zero[1, 0, 0] == 0.0 + def test_tiny_envelope_hessian(self) -> None: + """Both softmax paths keep the physical Hessian at a tiny envelope. + + The envelope enters the softmax through its logarithm, and in bridging + mode it carries the source gate, so the second derivative with respect + to it, which the force loss differentiates, must match the direct + physical form even where the envelope is far below one. + """ + from deepmd.dpmodel.descriptor.dpa4_nn.attention import ( + segment_envelope_gated_softmax as dp_softmax, + ) + from deepmd.pt.model.descriptor.sezm_nn.attention import ( + segment_envelope_gated_softmax as pt_softmax, + ) + + logits = torch.tensor( + [[[0.0]], [[20.0]]], dtype=torch.float64, device=PT_DEVICE + ) + dst = torch.zeros(2, dtype=torch.int64, device=PT_DEVICE) + z_bias_raw = torch.tensor( + [[math.log(math.expm1(1.0))]], dtype=torch.float64, device=PT_DEVICE + ) + eps = 1.0e-7 + null_mass = torch.nn.functional.softplus(z_bias_raw[0, 0]) + eps + + def physical_sum(edge_env: torch.Tensor) -> torch.Tensor: + edge_mass = edge_env**2 * torch.exp(logits[:, 0, 0]) + return (edge_mass / (null_mass + edge_mass.sum())).sum() + + def dp_sum(edge_env: torch.Tensor) -> torch.Tensor: + return dp_softmax(logits, edge_env[:, None], dst, 1, z_bias_raw, eps).sum() + + def pt_sum(edge_env: torch.Tensor) -> torch.Tensor: + return pt_softmax(logits, edge_env[:, None], dst, 1, z_bias_raw, eps).sum() + + edge_env = torch.tensor([1.0, 1.0e-15], dtype=torch.float64, device=PT_DEVICE) + reference = torch.autograd.functional.hessian(physical_sum, edge_env) + for attention_sum in (dp_sum, pt_sum): + hessian = torch.autograd.functional.hessian(attention_sum, edge_env) + assert bool(torch.isfinite(hessian).all()) + torch.testing.assert_close(hessian, reference, rtol=1.0e-9, atol=0.0) + def test_envelope_nextafter_cutoff_attention(self) -> None: """Adjacent float32 distances must remain stable in both implementations.""" from deepmd.dpmodel.descriptor.dpa4_nn.attention import ( @@ -2481,69 +2523,6 @@ def test_envelope_nextafter_cutoff_attention(self) -> None: ) assert 0.0 <= alpha_dp[1, 0, 0] < 1.0e-30 - def test_tiny_source_weight_hessian(self) -> None: - """The dpmodel and pt paths must preserve the physical Hessian.""" - from deepmd.dpmodel.descriptor.dpa4_nn.attention import ( - segment_envelope_gated_softmax as dp_softmax, - ) - from deepmd.pt.model.descriptor.sezm_nn.attention import ( - segment_envelope_gated_softmax as pt_softmax, - ) - - logits = torch.tensor( - [[[0.0]], [[20.0]]], dtype=torch.float32, device=PT_DEVICE - ) - edge_env = torch.ones((2, 1), dtype=torch.float32, device=PT_DEVICE) - dst = torch.zeros(2, dtype=torch.int64, device=PT_DEVICE) - z_bias_raw = torch.tensor( - [[math.log(math.expm1(1.0))]], dtype=torch.float32, device=PT_DEVICE - ) - eps = 1.0e-7 - - def dp_attention_sum(source_weight: torch.Tensor) -> torch.Tensor: - return dp_softmax( - logits, - edge_env, - dst, - 1, - z_bias_raw, - eps, - source_weight[:, None], - ).sum() - - def pt_attention_sum(source_weight: torch.Tensor) -> torch.Tensor: - return pt_softmax( - logits, - edge_env, - dst, - 1, - z_bias_raw, - eps, - source_weight[:, None], - ).sum() - - null_mass = torch.nn.functional.softplus(z_bias_raw[0, 0]) + eps - - def physical_sum(source_weight: torch.Tensor) -> torch.Tensor: - edge_mass = source_weight * torch.exp(logits[:, 0, 0]) - return (edge_mass / (null_mass + edge_mass.sum())).sum() - - source_weight = torch.tensor( - [1.0, 1.0e-30], dtype=torch.float32, device=PT_DEVICE - ) - hessian_dp = torch.autograd.functional.hessian(dp_attention_sum, source_weight) - hessian_pt = torch.autograd.functional.hessian(pt_attention_sum, source_weight) - reference = torch.autograd.functional.hessian(physical_sum, source_weight) - assert bool(torch.isfinite(hessian_dp).all()) - torch.testing.assert_close( - hessian_dp[0, 0], - reference[0, 0], - rtol=1.0e-5, - atol=1.0e-6, - ) - torch.testing.assert_close(hessian_dp, hessian_pt, rtol=1.0e-5, atol=32.0) - torch.testing.assert_close(hessian_dp, reference, rtol=1.0e-5, atol=32.0) - def test_segment_softmax_arbitrary_degree(self) -> None: # The destination scatter is layout-agnostic: E need not be a multiple # of n_nodes and dst may carry an arbitrary (non-row-major) order with a @@ -2672,8 +2651,8 @@ def test_so2_convolution_mlp_bias(self) -> None: pt_mod, dp_mod, kwargs = self._build_conv_pair(mlp_bias=True) self._assert_conv_parity(pt_mod, dp_mod, kwargs) - @pytest.mark.parametrize("n_atten_head", [0, 1]) # gate enters both paths - def test_so2_convolution_src_gate(self, n_atten_head) -> None: + @pytest.mark.parametrize("n_atten_head", [0, 1]) # gated envelope on both paths + def test_so2_convolution_gated_envelope(self, n_atten_head) -> None: pt_mod, dp_mod, kwargs = self._build_conv_pair(n_atten_head=n_atten_head) self._assert_conv_parity(pt_mod, dp_mod, kwargs, with_gate=True) @@ -2892,7 +2871,7 @@ def _build_gie_pair(self, lmax, channels): "zonal_provided", [False, True] ) # zonal_coupling: None (gather from Dt_full) vs provided-zonal input # path with D_node == D_cache only - @pytest.mark.parametrize("with_gate", [False, True]) # SFPG gate branch + @pytest.mark.parametrize("with_gate", [False, True]) # folded SFPG gate def test_gie(self, masked, zonal_provided, with_gate) -> None: lmax, channels = 2, 4 pt_mod, dp_mod = self._build_gie_pair(lmax, channels) @@ -3040,7 +3019,7 @@ def test_env_embedding(self, masked, mlp_bias) -> None: pt_mod, dp_mod = self._build_env_pair(mlp_bias=mlp_bias) self._assert_env_parity(pt_mod, dp_mod, masked=masked) - def test_env_embedding_src_gate(self) -> None: + def test_env_embedding_gated_envelope(self) -> None: pt_mod, dp_mod = self._build_env_pair() self._assert_env_parity(pt_mod, dp_mod, with_gate=True) @@ -3216,8 +3195,9 @@ def _dp_cache_from_padded( compute_dtype=np.float64, eps=eps, deg_norm_floor=deg_norm_floor, - inner_clamp=None, + bridging_clamp=None, bridging_switch=None, + edge_contact=None, edge_envelope=edge_envelope, radial_basis=radial_basis, random_gamma=random_gamma, @@ -3375,8 +3355,6 @@ def test_real_build_parity(self, local_nlist, deg_norm_floor) -> None: "inv_sqrt_deg", ): assert np.isfinite(np.asarray(getattr(dp_cache, name))).all(), name - # standard path carries no source gate - assert dp_cache.edge_src_gate is None assert dp_cache.D_to_m_cache == {} assert dp_cache.Dt_from_m_cache == {} @@ -3971,8 +3949,10 @@ def _assert_descr_parity(self, pt_mod, dp_mod, *, mapping=True) -> None: assert out_dp[0].shape == tuple(out_pt[0].shape) # descriptor-level tolerance: rtol 1e-10 / atol 1e-12 assert_parity(out_dp[0], out_pt[0], rtol=1e-10, atol=1e-12) - # unused returns are None on the dp side (pt returns empty tensors) - assert out_dp[1:] == (None, None, None, None) + # unused returns are None on the dp side (pt returns empty tensors); + # the trailing element is the source gate of a bridging window, absent + # from this unbridged fixture + assert out_dp[1:] == (None, None, None, None, None) @pytest.mark.parametrize("use_env_seed", [False, True]) # env FiLM + GIE seeding @pytest.mark.parametrize("n_blocks", [1, 2]) # interaction block stack depth @@ -4074,6 +4054,53 @@ def test_descriptor_no_mapping(self) -> None: pt_mod, dp_mod, _ = self._build_descr_pair() self._assert_descr_parity(pt_mod, dp_mod, mapping=False) + def test_descriptor_covalent_bridging_window(self) -> None: + """A covalent window sizes every pair by the elements at its ends. + + Both backends size the window from their own copy of the radius table, + so a multi-element ``type_map`` is what makes the per-pair radii differ + and the comparison meaningful. The window reaches the descriptor only + on the sparse-edge route, so the two are run on one shared ghost-free + graph. The unwindowed descriptor pins that the window actually moves + this fixture's output, so parity cannot hold vacuously. + """ + from deepmd.dpmodel.descriptor.dpa4 import ( + _graph_from_padded_nlist, + ) + + window = { + "type_map": ["O", "H", "N"], + "inner_clamp_f_inner": 1.0, + "inner_clamp_f_outer": 1.5, + "inner_clamp_scale": "covalent", + } + pt_mod, dp_mod, _ = self._build_descr_pair(**window) + np.testing.assert_array_equal( + dp_mod.contact_radius, pt_mod.contact_radius.cpu().numpy() + ) + _, plain_mod, _ = self._build_descr_pair() + + inp = self._inputs() + graph, atype_local = _graph_from_padded_nlist( + inp["coord"], inp["atype_ext"], inp["nlist"], inp["mapping"] + ) + out_dp, _, _ = dp_mod.call_graph(graph, atype_local) + out_pt, _, _, _ = pt_mod.forward_with_edges( + extended_coord=to_pt(inp["coord"][:, : self.nloc]), + extended_atype=to_pt(atype_local.reshape(self.nf, self.nloc)), + edge_index=to_pt(graph.edge_index), + edge_vec=to_pt(graph.edge_vec), + edge_mask=to_pt(graph.edge_mask), + ) + assert_parity( + np.reshape(out_dp, (self.nf, self.nloc, -1)), + out_pt, + rtol=1e-10, + atol=1e-12, + ) + out_plain, _, _ = plain_mod.call_graph(graph, atype_local) + assert np.max(np.abs(np.asarray(out_dp) - np.asarray(out_plain))) > 1e-8 + def test_descriptor_extra_node_l(self) -> None: # node degrees above message-passing degrees (GIE zonal wigner path) pt_mod, dp_mod, _ = self._build_descr_pair(extra_node_l=1) @@ -4249,7 +4276,7 @@ def _call(spin_arg): out_dp_s, out_pt_s = _call(spin) assert out_dp_s[0].shape == tuple(out_pt_s[0].shape) assert_parity(out_dp_s[0], out_pt_s[0], rtol=1e-10, atol=1e-12) - assert out_dp_s[1:] == (None, None, None, None) + assert out_dp_s[1:] == (None, None, None, None, None) # spin=None path: pt vs dp parity out_dp_n, out_pt_n = _call(None) @@ -4564,8 +4591,7 @@ class TestModelDefCompat: "type", "atomic_model", "bridging_method", - "bridging_r_inner", - "bridging_r_outer", + "bridging_window", "lora", } ) diff --git a/source/tests/pt/model/test_fitting_node_gate.py b/source/tests/pt/model/test_fitting_node_gate.py new file mode 100644 index 0000000000..f7805f2134 --- /dev/null +++ b/source/tests/pt/model/test_fitting_node_gate.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""The readout gate of a bridging window on the pt fitting networks. + +A bridged descriptor hands the fitting a per-atom source gate. The gate fades +the learned part of the output, the deviation of the output from the per-type +bias the fitting stores, so an atom whose gate is zero keeps that bias alone +and an atom whose gate is one is untouched. Both the generic energy fitting +and the SeZM one are covered, the latter on both of its branches: the plain +path and the one that conditions the network on a case embedding. +""" + +import numpy as np +import pytest +import torch + +from deepmd.pt.model.task.ener import ( + InvarFitting, +) +from deepmd.pt.model.task.sezm_ener import ( + SeZMEnergyFittingNet, +) +from deepmd.pt.utils import ( + env, +) + +from ...seed import ( + GLOBAL_SEED, +) + +NF, NLOC, ND, NTYPES = 2, 5, 8, 3 + + +def _build(kind: str, vacuum_ref: bool) -> torch.nn.Module: + """One fitting of each kind, in float64 on the test device.""" + if kind == "invar": + fitting = InvarFitting( + "energy", + NTYPES, + ND, + 1, + mixed_types=True, + precision="float64", + vacuum_ref=vacuum_ref, + seed=GLOBAL_SEED, + ) + else: + fitting = SeZMEnergyFittingNet( + NTYPES, + ND, + neuron=[8, 8], + precision="float64", + mixed_types=True, + dim_case_embd=2 if kind == "sezm_case_film" else 0, + case_film_embd=kind == "sezm_case_film", + vacuum_ref=vacuum_ref, + seed=GLOBAL_SEED, + ) + fitting = fitting.to(device=env.DEVICE, dtype=torch.float64) + if fitting.dim_case_embd > 0: + fitting.set_case_embd(0) + return fitting + + +def _inputs(vacuum_ref: bool): + rng = np.random.default_rng(GLOBAL_SEED) + descriptor = torch.as_tensor( + rng.normal(size=(NF, NLOC, ND)), dtype=torch.float64, device=env.DEVICE + ) + atype = torch.as_tensor(rng.integers(0, NTYPES, size=(NF, NLOC)), device=env.DEVICE) + node_gate = torch.as_tensor( + rng.uniform(0.0, 1.0, size=(NF, NLOC, 1)), + dtype=torch.float64, + device=env.DEVICE, + ) + kwargs = {} + if vacuum_ref: + kwargs["vacuum_descriptor"] = torch.as_tensor( + rng.normal(size=(NTYPES, ND)), dtype=torch.float64, device=env.DEVICE + ) + return descriptor, atype, node_gate, kwargs + + +def _constant_gate(value: float) -> torch.Tensor: + return torch.full((NF, NLOC, 1), value, dtype=torch.float64, device=env.DEVICE) + + +@pytest.mark.parametrize("kind", ["invar", "sezm", "sezm_case_film"]) +@pytest.mark.parametrize("vacuum_ref", [False, True]) +def test_the_gate_fades_the_deviation_from_the_bias( + kind: str, vacuum_ref: bool +) -> None: + """``out = bias + gate * (out_without_gate - bias)`` atom by atom.""" + fitting = _build(kind, vacuum_ref) + descriptor, atype, node_gate, kwargs = _inputs(vacuum_ref) + + ungated = fitting(descriptor, atype, **kwargs)["energy"] + gated = fitting(descriptor, atype, node_gate=node_gate, **kwargs)["energy"] + + reference = fitting.bias_atom_e[atype].to(torch.float64) + torch.testing.assert_close( + gated, + reference + node_gate * (ungated - reference), + rtol=0.0, + atol=1e-12, + ) + # a gate that changed nothing would make the comparison vacuous + assert (gated - ungated).abs().max().item() > 1e-6 + + +@pytest.mark.parametrize("kind", ["invar", "sezm", "sezm_case_film"]) +@pytest.mark.parametrize("vacuum_ref", [False, True]) +def test_the_two_ends_of_the_fade(kind: str, vacuum_ref: bool) -> None: + """A closed gate leaves the bias; an open one leaves the output alone.""" + fitting = _build(kind, vacuum_ref) + descriptor, atype, _, kwargs = _inputs(vacuum_ref) + + closed = fitting( + descriptor, + atype, + node_gate=_constant_gate(0.0), + **kwargs, + )["energy"] + torch.testing.assert_close( + closed, fitting.bias_atom_e[atype].to(torch.float64), rtol=0.0, atol=1e-12 + ) + + opened = fitting( + descriptor, + atype, + node_gate=_constant_gate(1.0), + **kwargs, + )["energy"] + torch.testing.assert_close( + opened, fitting(descriptor, atype, **kwargs)["energy"], rtol=0.0, atol=1e-12 + ) + + if vacuum_ref: + # the fold moves the reference out of the bias; the gate still reads + # the bias the atoms of a frozen pair fade to + bias = fitting.bias_atom_e.clone() + fitting.fold_vacuum_reference(kwargs["vacuum_descriptor"]) + assert fitting.folded_reference is not None + torch.testing.assert_close( + fitting.readout_reference(), bias, rtol=0.0, atol=1e-12 + ) + folded = fitting(descriptor, atype, node_gate=_constant_gate(0.0))["energy"] + torch.testing.assert_close(folded, closed, rtol=0.0, atol=1e-12) diff --git a/source/tests/pt/model/test_get_model_bridging.py b/source/tests/pt/model/test_get_model_bridging.py index 5c851f111d..614952641f 100644 --- a/source/tests/pt/model/test_get_model_bridging.py +++ b/source/tests/pt/model/test_get_model_bridging.py @@ -10,6 +10,9 @@ """ import copy +from pathlib import ( + Path, +) import pytest @@ -24,10 +27,7 @@ def _descriptor() -> dict: return { "type": "dpa4", "rcut": 4.0, - "rcut_smth": 0.5, "sel": 20, - "n_dim": 8, - "e_dim": 8, "precision": "float64", "seed": 7, } @@ -79,8 +79,12 @@ def test_canonical_builds_sezm_model() -> None: model = get_model(_canonical_config()) assert isinstance(model, SeZMModel) assert model.bridging_method == "ZBL" - assert model.bridging_r_inner == 0.8 - assert model.bridging_r_outer == 1.2 + # The window belongs to the descriptor, which is what the composition's + # radii have to reach for the learned side to honour them. + descriptor = model.atomic_model.descriptor + assert descriptor.bridging_f_inner == 0.8 + assert descriptor.bridging_f_outer == 1.2 + assert descriptor.bridging_scale == "absolute" def test_canonical_matches_sugar_serialize() -> None: @@ -252,6 +256,46 @@ def test_is_sezm_checkpoint_recognizes_canonical_params(tmp_path) -> None: assert is_sezm_checkpoint(ckpt2) +@pytest.mark.parametrize("method", ["none", "ZBL"]) # bridging disabled / enabled +def test_normalized_checkpoint_with_optional_radii_loads( + tmp_path: Path, method: str +) -> None: + """Checkpoint loading preserves the normalized optional window radii.""" + import numpy as np + import torch + + from deepmd.infer import ( + DeepPot, + ) + from deepmd.pt.train.wrapper import ( + ModelWrapper, + ) + from deepmd.utils.argcheck import ( + model_args, + ) + + config = _sugar_config() + config["bridging_method"] = method + config.pop("bridging_r_inner") + config.pop("bridging_r_outer") + config["descriptor"].update(channels=4, n_radial=3, lmax=1, mmax=1, n_blocks=1) + config["use_compile"] = False + params = model_args().normalize_value(config, trim_pattern="_*") + model_args().check_value(params, strict=True) + assert params["bridging_r_inner"] is None + assert params["bridging_r_outer"] is None + model = get_model(copy.deepcopy(params)).eval() + wrapper = ModelWrapper(model, model_params=params) + checkpoint = tmp_path / "model.pt" + torch.save({"model": wrapper.state_dict()}, checkpoint) + + evaluator = DeepPot(str(checkpoint)) + assert evaluator.get_type_map() == params["type_map"] + results = evaluator.eval([[0.0, 0.0, 0.0, 0.9, 0.0, 0.0]], None, [0, 1]) + assert all(np.isfinite(value).all() for value in results) + evaluator.close() + + def test_sugar_with_top_level_lora_builds() -> None: """The concise dpa4+bridging form with trainer-owned top-level `lora` must keep building: the expansion routes `lora` to the composition diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py index ef0beb3c68..18ca49b8bf 100644 --- a/source/tests/pt/model/test_sezm_model.py +++ b/source/tests/pt/model/test_sezm_model.py @@ -708,8 +708,11 @@ def test_fixed_edge_geometry_matches_standard_cache(self) -> None: compute_dtype=descriptor.compute_dtype, eps=descriptor.eps, deg_norm_floor=descriptor.deg_norm_floor, - inner_clamp=descriptor.inner_clamp, + bridging_clamp=descriptor.bridging_clamp, bridging_switch=descriptor.bridging_switch, + # The standard path this compares against is bridging-free, so the + # fixture configures no window and no pair length scale is needed. + edge_contact=None, edge_envelope=descriptor.edge_envelope, radial_basis=descriptor.radial_basis, has_exclude_types=False, @@ -1346,7 +1349,7 @@ def test_compile_matches_eager_and_backpropagates(self) -> None: class TestInnerPotential(unittest.TestCase): - """Test InnerPotential ZBL analytical pair potential.""" + """Test InnerPotential analytical pair potential.""" def setUp(self) -> None: self.device = env.DEVICE @@ -1367,6 +1370,48 @@ def _pair_edges( edge_mask = torch.tensor([True, True], device=self.device) return edge_vec, edge_index, atype_flat, edge_mask + def test_nlh_matches_the_backend_agnostic_reference(self) -> None: + """The pt forward reproduces the dpmodel one in NLH mode.""" + from deepmd.dpmodel.atomic_model.inner_potential import ( + InnerPotential as InnerPotentialDP, + ) + + type_map = ["O", "H", "Fe"] + pot = InnerPotential(type_map=type_map, mode="nlh").to(self.device) + reference = InnerPotentialDP(type_map=type_map, mode="nlh") + for r, pair in ((0.4, [0, 1]), (1.1, [0, 0]), (2.5, [1, 2])): + edge_vec, edge_index, atype_flat, edge_mask = self._pair_edges(r, pair) + got = pot(edge_vec, edge_index, atype_flat, edge_mask, 2) + want = reference.call( + edge_vec.detach().cpu().numpy(), + edge_index.detach().cpu().numpy(), + atype_flat.detach().cpu().numpy(), + edge_mask.detach().cpu().numpy(), + 2, + ) + np.testing.assert_allclose( + got.detach().cpu().numpy(), np.asarray(want), rtol=1e-12 + ) + + def test_nlh_differs_from_zbl(self) -> None: + """The mode reaches the energy rather than falling back to ZBL.""" + edge_vec, edge_index, atype_flat, edge_mask = self._pair_edges(0.8, [0, 0]) + values = [ + float( + InnerPotential(type_map=["O"], mode=mode) + .to(self.device)(edge_vec, edge_index, atype_flat, edge_mask, 2) + .sum() + ) + for mode in ("zbl", "nlh") + ] + self.assertGreater(values[1], 0.0) + self.assertGreater(abs(values[1] / values[0] - 1.0), 0.05) + + def test_nlh_series_table_stays_out_of_the_checkpoint(self) -> None: + """Every array follows type_map and mode, so none of them is stored.""" + pot = InnerPotential(type_map=["O", "H"], mode="nlh").to(self.device) + self.assertEqual(list(pot.state_dict()), []) + def test_zbl_known_value_OO(self) -> None: """ZBL energy for an O-O pair matches the analytic reference.""" pot = InnerPotential(type_map=["O", "H"], mode="ZBL").to(self.device) @@ -1422,36 +1467,6 @@ def test_zbl_gradient_exists(self) -> None: self.assertIsNotNone(edge_vec.grad) self.assertTrue(torch.isfinite(edge_vec.grad).all()) - def test_virtual_spin_types_masked(self) -> None: - """Edges touching a virtual spin type (>= real_type_count) contribute 0.""" - pot = InnerPotential(type_map=["O", "H"], mode="ZBL").to(self.device) - # Node 2 is a virtual spin atom (type 2 >= real_type_count=2). - edge_vec = torch.tensor( - [[1.0, 0.0, 0.0], [-1.0, 0.0, 0.0], [0.5, 0.0, 0.0], [-0.5, 0.0, 0.0]], - dtype=torch.float64, - device=self.device, - ) - # Edges: (0<->1) real-real, (0<->2) touch virtual node 2. - edge_index = torch.tensor( - [[1, 0, 2, 0], [0, 1, 0, 2]], dtype=torch.long, device=self.device - ) - atype_flat = torch.tensor([0, 1, 2], dtype=torch.long, device=self.device) - edge_mask = torch.tensor([True, True, True, True], device=self.device) - - with_virtual = pot( - edge_vec, edge_index, atype_flat, edge_mask, n_node=3, real_type_count=2 - ) - # Only the real-real pair survives. - real_only = pot( - edge_vec[:2], - edge_index[:, :2], - atype_flat, - edge_mask[:2], - n_node=3, - real_type_count=2, - ) - torch.testing.assert_close(with_virtual, real_only) - def test_unknown_element_raises(self) -> None: """Test that unknown element raises ValueError.""" with self.assertRaises(ValueError): @@ -2213,7 +2228,6 @@ def _build_model_params(self, *, bridging_method: str = "none") -> dict: "rcut": 3.0, "channels": 4, "n_focus": 1, - "focus_compete": False, "n_radial": 3, "radial_mlp": [6], "use_env_seed": False, @@ -2256,7 +2270,43 @@ def test_bridging_zbl_creates_potential(self) -> None: model = get_sezm_model(self._build_model_params(bridging_method="ZBL")) self.assertIsNotNone(model.inter_potential) self.assertEqual(model.bridging_method, "ZBL") - self.assertIsNotNone(model.atomic_model.descriptor.inner_clamp) + self.assertIsNotNone(model.atomic_model.descriptor.bridging_clamp) + + def test_legacy_bridged_state_is_rejected(self) -> None: + """Tagged and untagged checkpoints obey the descriptor version contract.""" + model = get_sezm_model(self._build_model_params(bridging_method="ZBL")) + for version in (None, 1.0, 1.1, 1.2): + with self.subTest(version=version): + state = model.state_dict() + version_key = "atomic_model.descriptor.version_tensor" + if version is None: + del state[version_key] + else: + state[version_key] = torch.full_like(state[version_key], version) + with self.assertRaisesRegex(ValueError, "Retrain the bridged model"): + model.load_state_dict(state) + + def test_legacy_plain_state_into_bridged_model_requires_finetune(self) -> None: + """Direct loading rejects legacy plain weights without misidentifying them.""" + plain = get_sezm_model(self._build_model_params(bridging_method="none")) + bridged = get_sezm_model(self._build_model_params(bridging_method="ZBL")) + state = plain.state_dict() + version_key = "atomic_model.descriptor.version_tensor" + state[version_key] = torch.full_like(state[version_key], 1.2) + with self.assertRaisesRegex( + ValueError, + r"format version 1\.2 into a bridged descriptor.*--finetune", + ): + bridged.load_state_dict(state, strict=False) + + def test_current_bridged_state_round_trips(self) -> None: + """The state-dict and portable representations retain the same tensors.""" + params = self._build_model_params(bridging_method="ZBL") + model = get_sezm_model(params) + restored = get_sezm_model(params) + restored.load_state_dict(model.state_dict()) + portable = SeZMModel.deserialize(restored.serialize()) + torch.testing.assert_close(portable.state_dict(), model.state_dict()) def test_empty_statistics_raise(self) -> None: """A fully filtered statistics sample cannot calibrate output bias.""" @@ -2445,41 +2495,99 @@ def test_set_by_statistic_fits_raw_labels_by_definition(self) -> None: repeated_bias = model.get_out_bias().detach().cpu().numpy().reshape(-1)[:2] np.testing.assert_allclose(repeated_bias, raw_fit, atol=1.0e-8) + def _excluded_pair_models( + self, exclusion_key: str, exclusion_value: list + ) -> tuple[SeZMModel, SeZMModel]: + """A plain and a bridged float64 model carrying the same weights.""" + plain_params = self._build_model_params(bridging_method="none") + zbl_params = self._build_model_params(bridging_method="ZBL") + for params in (plain_params, zbl_params): + params[exclusion_key] = exclusion_value + params["descriptor"]["precision"] = "float64" + params["fitting_net"]["precision"] = "float64" + model_plain = get_sezm_model(plain_params).to(self.device).eval() + model_zbl = get_sezm_model(zbl_params).to(self.device).eval() + model_zbl.load_state_dict(model_plain.state_dict(), strict=False) + return model_plain, model_zbl + def test_zbl_respects_exclusions(self) -> None: - """Excluded atoms and pairs contribute neither learned nor ZBL energy.""" - coord = torch.tensor( - [[[0.0, 0.0, 0.0], [0.8, 0.0, 0.0]]], - dtype=torch.float64, - device=self.device, - ) + """Excluded pairs and atoms leave the analytical term as they leave the learned one. + + The two exclusions differ in what they hide. A pair exclusion takes the + pair out of the descriptor and out of the analytical term together, so + a bridged model reproduces a plain one even with the two atoms at the + inner radius of the window. An atom exclusion acts on atomic outputs + alone: the excluded atom loses its half of the pair energy with the + rest of its output, its geometry still shapes the environment of its + partner, and the partner keeps its own half. Beyond the outer radius + the bridged model is therefore the plain one plus that half; at the + inner radius the partner is a frozen atom and is left with the per-type + bias of the fitting plus the same half. + """ atype = torch.tensor([[0, 1]], dtype=torch.long, device=self.device) box = torch.tensor( [[10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0]], dtype=torch.float64, device=self.device, ) - for exclusion_key, exclusion_value in ( - ("pair_exclude_types", [[0, 1]]), - ("atom_exclude_types", [1]), - ): - with self.subTest(exclusion_key=exclusion_key): - plain_params = self._build_model_params(bridging_method="none") - zbl_params = self._build_model_params(bridging_method="ZBL") - for params in (plain_params, zbl_params): - params[exclusion_key] = exclusion_value - params["descriptor"]["precision"] = "float64" - params["fitting_net"]["precision"] = "float64" - - model_plain = get_sezm_model(plain_params).to(self.device).eval() - model_zbl = get_sezm_model(zbl_params).to(self.device).eval() - model_zbl.load_state_dict(model_plain.state_dict(), strict=False) - torch.testing.assert_close( - model_zbl(coord, atype, box=box)["energy"], - model_plain(coord, atype, box=box)["energy"], - atol=1.0e-12, - rtol=1.0e-12, - ) + def _pair(separation: float) -> torch.Tensor: + return torch.tensor( + [[[0.0, 0.0, 0.0], [separation, 0.0, 0.0]]], + dtype=torch.float64, + device=self.device, + ) + + def _half_pair_energy(model: SeZMModel, separation: float) -> torch.Tensor: + """The analytical half-energy each atom of the pair collects.""" + edge_vec = torch.tensor( + [[separation, 0.0, 0.0], [-separation, 0.0, 0.0]], + dtype=torch.float64, + device=self.device, + ) + edge_index = torch.tensor([[1, 0], [0, 1]], device=self.device) + return model.inter_potential( + edge_vec, + edge_index, + atype.reshape(-1), + torch.ones(2, dtype=torch.bool, device=self.device), + 2, + ).reshape(-1)[0] + + with self.subTest(exclusion_key="pair_exclude_types"): + plain, bridged = self._excluded_pair_models("pair_exclude_types", [[0, 1]]) + coord = _pair(0.8) + torch.testing.assert_close( + bridged(coord, atype, box=box)["energy"], + plain(coord, atype, box=box)["energy"], + atol=1.0e-12, + rtol=1.0e-12, + ) + + with self.subTest(exclusion_key="atom_exclude_types"): + plain, bridged = self._excluded_pair_models("atom_exclude_types", [1]) + coord = _pair(1.5) # beyond the outer radius: the window is shut + torch.testing.assert_close( + bridged(coord, atype, box=box)["energy"], + plain(coord, atype, box=box)["energy"] + + _half_pair_energy(bridged, 1.5), + atol=1.0e-12, + rtol=1.0e-12, + ) + + coord = _pair(0.8) # at the inner radius: the pair is frozen + bias = bridged.atomic_model.fitting_net.bias_atom_e.reshape(-1) + torch.testing.assert_close( + bridged(coord, atype, box=box)["atom_energy"].reshape(-1), + torch.stack( + [ + bias[0] + _half_pair_energy(bridged, 0.8), + torch.zeros_like(bias[0]), + ] + ), + atol=1.0e-12, + rtol=1.0e-12, + ) class TestSeZMPhantomAtoms(unittest.TestCase): @@ -2506,7 +2614,6 @@ def _build_model(self) -> SeZMModel: "rcut": 4.0, "channels": 4, "n_focus": 1, - "focus_compete": False, "n_radial": 3, "radial_mlp": [6], "use_env_seed": False, @@ -2645,7 +2752,6 @@ def _build_model_params( "rcut": 3.0, "channels": 4, "n_focus": 1, - "focus_compete": False, "n_radial": 3, "radial_mlp": [6], "use_env_seed": False, diff --git a/source/tests/pt/model/test_sezm_nvalchemi.py b/source/tests/pt/model/test_sezm_nvalchemi.py index 836d46524e..8030040bc6 100644 --- a/source/tests/pt/model/test_sezm_nvalchemi.py +++ b/source/tests/pt/model/test_sezm_nvalchemi.py @@ -24,12 +24,12 @@ import torch +from deepmd.dpmodel.atomic_model.inner_potential import ( + ELEMENT_TO_Z, +) from deepmd.pt.model.model import ( get_sezm_model, ) -from deepmd.pt.model.model.sezm_model import ( - ELEMENT_TO_Z, -) from deepmd.pt.utils import ( env, ) diff --git a/source/tests/pt/model/test_sezm_parallel.py b/source/tests/pt/model/test_sezm_parallel.py index bceec317d8..9be76a22b6 100644 --- a/source/tests/pt/model/test_sezm_parallel.py +++ b/source/tests/pt/model/test_sezm_parallel.py @@ -336,7 +336,7 @@ def test_descriptor_parity_cpu(self) -> None: sysm = _build_extended_system(model, device) comm = _self_comm_dict(sysm["mapping"], sysm["nloc"], sysm["nall"]) - ref, _, ref_vacuum = descriptor.forward_with_edges( + ref, _, ref_vacuum, _ = descriptor.forward_with_edges( extended_coord=sysm["coord"][:, : sysm["nloc"], :], extended_atype=sysm["atype"], edge_index=sysm["edge_index"], @@ -344,7 +344,7 @@ def test_descriptor_parity_cpu(self) -> None: edge_mask=sysm["edge_mask"], vacuum_conditions={}, ) - par, _, par_vacuum = descriptor.forward_with_edges( + par, _, par_vacuum, _ = descriptor.forward_with_edges( extended_coord=sysm["coord"], extended_atype=sysm["extended_atype"], edge_index=sysm["edge_scatter_index"], diff --git a/source/tests/pt/model/test_sezm_parallel_bridging_parity.py b/source/tests/pt/model/test_sezm_parallel_bridging_parity.py index cabde13078..04d9eae8ef 100644 --- a/source/tests/pt/model/test_sezm_parallel_bridging_parity.py +++ b/source/tests/pt/model/test_sezm_parallel_bridging_parity.py @@ -2,11 +2,13 @@ """Issue #5906 Task 2 (pt backend): SFPG completion across ranks. A BRIDGED SeZM model's parallel (comm_dict) path must reproduce the folded -single-domain path. The Source Freeze Propagation Gate folds each node's -full outgoing-edge set; under domain decomposition a rank only holds edges -with owned destinations, so the src-keyed per-node partials are incomplete -and must be completed by one reverse-accumulate + forward-broadcast border -exchange before the gate is applied. +single-domain path. The Source Freeze Propagation Gate multiplies, into the +envelope of an edge, the switching amplitudes of every pair of that edge's +source other than its own, and hands the full per-node product to the model; +under domain decomposition a rank only holds edges with owned destinations, +so the src-keyed per-node partials are incomplete and must be completed by +one reverse-accumulate + forward-broadcast border exchange before either +gate is applied. Geometry is the load-bearing part: a sub-``r_outer`` pair STRADDLES the periodic x boundary so the close contact's src is a ghost image row. @@ -171,8 +173,15 @@ def _assert_parity(self, gap: float) -> None: msg="reference forces are ~0; the parity check would be vacuous", ) for key in ("energy", "extended_force", "virial"): + # The perturbed weights give outputs of order 1e6 to 1e7, and the + # two paths differ by round-off of about 1e-14 of the largest entry, + # so the tolerance follows that entry rather than each one. torch.testing.assert_close( - par[key], ref[key], rtol=1e-8, atol=1e-9, msg=f"mismatch in {key}" + par[key], + ref[key], + rtol=0.0, + atol=1e-11 * ref[key].abs().max().item(), + msg=f"mismatch in {key}", ) def test_parity_hard_freeze_pair_cpu(self) -> None: diff --git a/source/tests/pt/model/test_sezm_spin_model.py b/source/tests/pt/model/test_sezm_spin_model.py index 314a715364..ffb86609fb 100644 --- a/source/tests/pt/model/test_sezm_spin_model.py +++ b/source/tests/pt/model/test_sezm_spin_model.py @@ -366,153 +366,15 @@ def test_dens_mode_is_rejected(self) -> None: with self.assertRaises(NotImplementedError): prepare_model_for_loss(model, {"type": "dens"}) - def test_bridging_masks_virtual_pairs(self) -> None: - """ZBL bridging should ignore virtual spin types without indexing them.""" - model = get_model(self._build_model_params(bridging_method="ZBL")).to( - self.device - ) - self.assertIsNotNone(model.inter_potential) - - # Nodes 0 (type 0) and 1 (type 1) are real; node 2 (type 2) is a virtual - # spin atom (>= real_type_count=2). Edges touching node 2 must be masked. - edge_vec = torch.tensor( - [[1.0, 0.0, 0.0], [-1.0, 0.0, 0.0], [0.5, 0.0, 0.0], [-0.5, 0.0, 0.0]], - dtype=torch.float64, - device=self.device, - ) - edge_index = torch.tensor( - [[1, 0, 2, 0], [0, 1, 0, 2]], dtype=torch.long, device=self.device - ) - atype_flat = torch.tensor([0, 1, 2], dtype=torch.long, device=self.device) - edge_mask = torch.tensor([True, True, True, True], device=self.device) - - energy_with_virtual = model.inter_potential( - edge_vec, edge_index, atype_flat, edge_mask, n_node=3, real_type_count=2 - ) - energy_real_only = model.inter_potential( - edge_vec[:2], - edge_index[:, :2], - atype_flat, - edge_mask[:2], - n_node=3, - real_type_count=2, - ) - - torch.testing.assert_close(energy_with_virtual, energy_real_only) - - def test_zbl_change_out_bias_is_invariant_for_self_labels(self) -> None: - """Spin-expanded statistics include the complete bridged energy.""" - model = get_model(self._build_model_params(bridging_method="ZBL")).to( - self.device - ) - model.eval() - label = model( - self.coord, - self.atype, - spin=self.spin, - box=self.box, - )["energy"].detach() - old_bias = model.get_out_bias().detach().clone() - sample = { - "coord": self.coord, - "atype": self.atype, - "spin": self.spin, - "box": self.box, - "energy": label, - "natoms": torch.tensor( - [[3, 3, 2, 1]], - dtype=torch.long, - device=self.device, - ), - "find_energy": torch.tensor(1.0, device=self.device), - } - - model.change_out_bias( - [sample], - bias_adjust_mode="change-by-statistic", - ) - torch.testing.assert_close( - model.get_out_bias(), - old_bias, - atol=1.0e-12, - rtol=1.0e-12, - ) - - def test_zbl_statistics_use_physical_spin_neighbor_topology(self) -> None: - """Statistics preserve real-nlist expansion when `sel` is saturated.""" - params = self._build_model_params(bridging_method="ZBL") - params["descriptor"]["sel"] = [1, 1] - params["descriptor"]["precision"] = "float64" - params["fitting_net"]["precision"] = "float64" - model = get_model(params).to(self.device).eval() - coord = torch.tensor( - [ - [ - [0.0, 0.0, 0.0], - [0.8, 0.0, 0.0], - [0.0, 0.9, 0.0], - [0.9, 0.9, 0.0], - ] - ], - dtype=torch.float64, - device=self.device, - ) - atype = torch.tensor( - [[0, 1, 0, 1]], - dtype=torch.long, - device=self.device, - ) - spin = torch.tensor( - [ - [ - [0.20, 0.10, 0.00], - [0.00, 0.00, 0.00], - [0.10, 0.20, 0.10], - [0.00, 0.00, 0.00], - ] - ], - dtype=torch.float64, - device=self.device, - ) - box = torch.tensor( - [[6.0, 0.0, 0.0, 0.0, 6.0, 0.0, 0.0, 0.0, 6.0]], - dtype=torch.float64, - device=self.device, - ) - label = model(coord, atype, spin=spin, box=box)["energy"].detach() - stat_energy = model.predict_atomic_outputs_for_stat( - coord, - atype, - box, - spin=spin, - )["energy"].sum(dim=1) - torch.testing.assert_close(stat_energy, label, atol=1.0e-12, rtol=1.0e-12) - - old_bias = model.get_out_bias().detach().clone() - model.change_out_bias( - [ - { - "coord": coord, - "atype": atype, - "spin": spin, - "box": box, - "energy": label, - "natoms": torch.tensor( - [[4, 4, 2, 2]], - dtype=torch.long, - device=self.device, - ), - "find_energy": torch.tensor(1.0, device=self.device), - } - ], - bias_adjust_mode="change-by-statistic", - ) - torch.testing.assert_close( - model.get_out_bias(), - old_bias, - atol=1.0e-12, - rtol=1.0e-12, - ) + def test_bridging_is_rejected(self) -> None: + """The virtual-atom scheme does not compose with an analytical potential. + + A magnetic atom sits within the bridging window of its own virtual + partner, so a bridged descriptor would freeze every magnetic atom; the + model refuses the combination instead of building it. + """ + with self.assertRaisesRegex(ValueError, "native"): + get_model(self._build_model_params(bridging_method="ZBL")) @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON) def test_compile_matches_eager(self) -> None: diff --git a/source/tests/pt/model/test_sezm_vacuum_freeze.py b/source/tests/pt/model/test_sezm_vacuum_freeze.py index 4a2990aa9a..dc3a9766ec 100644 --- a/source/tests/pt/model/test_sezm_vacuum_freeze.py +++ b/source/tests/pt/model/test_sezm_vacuum_freeze.py @@ -43,6 +43,15 @@ class TestSeZMVacuumFreeze(unittest.TestCase): def test_frozen_model_folds_the_reference(self) -> None: self.check_frozen_model(numb_fparam=0) + def test_frozen_bridged_model_keeps_the_readout_reference(self) -> None: + """The readout gate of a bridged model fades to the bias the fold moved. + + A pair inside the bridging window keeps the gate open in part, so the + frozen archive reproduces the eager energies only if the gate reads + the isolated-atom output rather than the folded bias. + """ + self.check_frozen_model(numb_fparam=0, bridged=True) + def test_frozen_model_keeps_the_reference_with_fparam(self) -> None: """With frame parameters the archive references from the stored vacuum table.""" self.check_frozen_model(numb_fparam=1) @@ -67,9 +76,21 @@ def test_dens_checkpoint_is_rejected(self) -> None: with self.assertRaisesRegex(ValueError, "`ener` mode"): freeze_sezm_to_pt2(str(ckpt), str(Path(tmp) / "dens.pt2")) - def check_frozen_model(self, numb_fparam: int, device: torch.device = _CPU) -> None: + def check_frozen_model( + self, + numb_fparam: int, + device: torch.device = _CPU, + bridged: bool = False, + ) -> None: params = _tiny_sezm_model_params() - params["preset_out_bias"] = {"energy": {"A": BIAS[0, 0], "B": BIAS[1, 0]}} + if bridged: + # the analytical term and the covalent window need element names + params["type_map"] = ["Ni", "O"] + params["bridging_method"] = "zbl" + types = params["type_map"] + params["preset_out_bias"] = { + "energy": {types[0]: BIAS[0, 0], types[1]: BIAS[1, 0]} + } params["fitting_net"]["vacuum_ref"] = True params["fitting_net"]["numb_fparam"] = numb_fparam fparam = None if numb_fparam == 0 else np.array([[0.7]]) @@ -89,6 +110,9 @@ def check_frozen_model(self, numb_fparam: int, device: torch.device = _CPU) -> N natoms = 5 atype = np.array([0, 1, 0, 1, 0], dtype=np.int32) coord = rng.random((1, natoms, 3)) * box_edge * 0.4 + box_edge * 0.3 + if bridged: + # a Ni-Ni pair inside its window, where the readout gate is partly open + coord[0, 2] = coord[0, 0] + np.array([1.0, 0.0, 0.0]) eager = ( model.forward( torch.tensor(coord, dtype=torch.float64, device=device), diff --git a/source/tests/pt/test_finetune.py b/source/tests/pt/test_finetune.py index 8998392830..de3ecad55a 100644 --- a/source/tests/pt/test_finetune.py +++ b/source/tests/pt/test_finetune.py @@ -367,5 +367,146 @@ def setUp(self) -> None: self.testkey = None +model_sezm = { + "type": "dpa4", + "type_map": ["O", "H", "B"], + "descriptor": { + "type": "dpa4", + "rcut": 4.0, + "sel": 20, + "precision": "float64", + "seed": 7, + }, + "fitting_net": { + "type": "dpa4_ener", + "neuron": [4, 4], + "precision": "float64", + "seed": 7, + }, +} + + +class TestSeZMBridgingFinetune(unittest.TestCase): + """A DPA4 checkpoint fine-tunes into its ZBL-bridged form and back. + + The pt backend builds the analytical term inside ``SeZMModel``, so both + models hold the same trained tensors and differ by the nuclear charges of + the bridging term alone, which the type map fixes. + """ + + def setUp(self) -> None: + self.tmpdir = tempfile.mkdtemp(prefix="pt_sezm_bridging_") + self.cwd = os.getcwd() + os.chdir(self.tmpdir) + + def tearDown(self) -> None: + os.chdir(self.cwd) + shutil.rmtree(self.tmpdir, ignore_errors=True) + + @staticmethod + def _config(model: dict) -> dict: + input_json = str(Path(__file__).parent / "water/se_atten.json") + with open(input_json) as f: + config = json.load(f) + data_file = [str(Path(__file__).parent / "water/data/single")] + config["training"]["training_data"]["systems"] = data_file + config["training"]["validation_data"]["systems"] = data_file + config["training"]["numb_steps"] = 1 + config["training"]["save_freq"] = 1 + config["model"] = deepcopy(model) + return config + + def _finetune_across_bridging( + self, source: dict, target: dict, source_version: float | None = None + ) -> torch.nn.Module: + config = self._config(source) + trainer = get_trainer(config) + if source_version is not None: + descriptor = trainer.model.get_descriptor() + descriptor.version = source_version + descriptor.version_tensor.fill_(source_version) + trainer.run() + checkpoint = config["training"].get("save_ckpt", "model.ckpt") + ".pt" + + finetune_config = self._config(target) + finetune_config["model"], finetune_links = get_finetune_rules( + checkpoint, finetune_config["model"] + ) + trainer_finetune = get_trainer( + finetune_config, + finetune_model=checkpoint, + finetune_links=finetune_links, + ) + + pretrained = torch.load(checkpoint, map_location=env.DEVICE, weights_only=True)[ + "model" + ] + state = trainer_finetune.wrapper.state_dict() + learned = [ + key for key in state if ".descriptor." in key or ".fitting_net." in key + ] + self.assertTrue(learned) + for key in learned: + if source_version is not None and key.endswith(".version_tensor"): + continue + torch.testing.assert_close( + state[key], + pretrained[key], + msg=f"{key} should come from the checkpoint", + ) + return trainer_finetune.model + + def test_finetune_plain_into_bridged(self) -> None: + bridged = deepcopy(model_sezm) + bridged.update( + bridging_method="ZBL", bridging_r_inner=0.5, bridging_r_outer=0.8 + ) + self._finetune_across_bridging(model_sezm, bridged) + + def test_finetune_bridged_into_plain(self) -> None: + bridged = deepcopy(model_sezm) + bridged.update( + bridging_method="ZBL", bridging_r_inner=0.5, bridging_r_outer=0.8 + ) + self._finetune_across_bridging(bridged, model_sezm) + + def test_legacy_plain_finetune_has_a_portable_bridged_record(self) -> None: + """A version-1.2 plain checkpoint acquires the current bridging semantics.""" + from deepmd.dpmodel.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4DP + from deepmd.pt_expt.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4PTExpt + + bridged = deepcopy(model_sezm) + bridged.update( + bridging_method="ZBL", bridging_r_inner=0.5, bridging_r_outer=0.8 + ) + model = self._finetune_across_bridging(model_sezm, bridged, source_version=1.2) + model.eval() + restored = type(model).deserialize(model.serialize()).to(env.DEVICE).eval() + self.assertEqual( + restored.get_descriptor().version, model.get_descriptor().LATEST_VERSION + ) + descriptor_record = model.get_descriptor().serialize() + for descriptor_cls in (DescrptDPA4DP, DescrptDPA4PTExpt): + converted = descriptor_cls.deserialize(deepcopy(descriptor_record)) + self.assertEqual(converted.serialize()["@version"], 1.3) + inputs = { + "coord": torch.tensor( + [[[0.0, 0.0, 0.0], [0.6, 0.0, 0.0], [2.0, 0.7, 0.0]]], + dtype=torch.float64, + device=env.DEVICE, + ), + "atype": torch.tensor([[0, 1, 1]], device=env.DEVICE), + "box": (6.0 * torch.eye(3, dtype=torch.float64, device=env.DEVICE)).reshape( + 1, 9 + ), + } + expected = model(**inputs) + actual = restored(**inputs) + for key in ("energy", "force", "virial"): + torch.testing.assert_close( + actual[key], expected[key], rtol=1e-12, atol=1e-12 + ) + + if __name__ == "__main__": unittest.main() diff --git a/source/tests/pt/test_training.py b/source/tests/pt/test_training.py index 9cce8c02aa..85a8a1a749 100644 --- a/source/tests/pt/test_training.py +++ b/source/tests/pt/test_training.py @@ -111,6 +111,10 @@ class TestStatisticsFrameFiltering(unittest.TestCase): """Verify statistics use the training minimum-distance frame semantics.""" def test_min_pair_dist_filters_every_frame_aligned_tensor(self) -> None: + # The field holds margins against each frame's own pair thresholds, so + # the verdict is a comparison against one. Both margins clear the + # configured window radius of 0.8, which is therefore no threshold of + # its own: a filter that compared against it would keep both frames. batch = { "coord": torch.arange(18, dtype=torch.float64, device="cpu").reshape( 2, 3, 3 @@ -122,9 +126,10 @@ def test_min_pair_dist_filters_every_frame_aligned_tensor(self) -> None: [[1000.0], [2.0]], dtype=torch.float64, device="cpu" ), "natoms": torch.tensor([[3, 3, 2, 1], [3, 3, 1, 2]], device="cpu"), - "min_pair_dist": torch.tensor( - [[0.5], [1.0]], dtype=torch.float64, device="cpu" + "pair_margin": torch.tensor( + [[0.9], [1.2]], dtype=torch.float64, device="cpu" ), + "find_pair_margin": np.float32(1.0), "fid": ["rejected", "accepted"], "find_energy": np.float32(1.0), } @@ -149,12 +154,16 @@ def test_min_pair_dist_filters_every_frame_aligned_tensor(self) -> None: ) def test_min_pair_dist_skips_system_without_valid_frames(self) -> None: + # The only frame falls short of a margin of one while still clearing + # the configured window radius, so a filter reading that radius as a + # threshold would keep the system instead of skipping it. batch = { "coord": torch.zeros(1, 2, 3, device="cpu"), "atype": torch.zeros(1, 2, dtype=torch.long, device="cpu"), "energy": torch.zeros(1, 1, device="cpu"), "natoms": torch.tensor([[2, 2, 2]], device="cpu"), - "min_pair_dist": torch.tensor([[0.5]], device="cpu"), + "pair_margin": torch.tensor([[0.9]], device="cpu"), + "find_pair_margin": np.float32(1.0), "find_energy": np.float32(1.0), } sampled = make_stat_input( @@ -171,13 +180,21 @@ def test_distributed_filter_skips_when_any_rank_is_empty(self) -> None: def mark_remote_rank_empty( valid_flag: torch.Tensor, op: Any, + group: Any, ) -> None: self.assertEqual(op, torch.distributed.ReduceOp.MIN) + self.assertFalse(valid_flag.is_cuda) valid_flag.zero_() - with patch( - "deepmd.pt.train.training.dist.all_reduce", - side_effect=mark_remote_rank_empty, + # The agreement travels over a dedicated CPU process group, which only + # exists once distributed training is initialized; stubbing the lookup + # keeps the MIN decision itself under test without a live group. + with ( + patch( + "deepmd.pt.train.training.dist.all_reduce", + side_effect=mark_remote_rank_empty, + ), + patch("deepmd.pt.train.training._validity_group", return_value=None), ): self.assertFalse(all_ranks_have_valid_frames(local_has_valid=True)) diff --git a/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py b/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py index a72698fbd4..de05319045 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py +++ b/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py @@ -81,8 +81,9 @@ def _make_descriptor( random_gamma=False, precision=precision, seed=7, - inner_clamp_r_inner=0.8 if source_gated else None, - inner_clamp_r_outer=1.2 if source_gated else None, + inner_clamp_f_inner=0.8 if source_gated else None, + inner_clamp_f_outer=1.2 if source_gated else None, + inner_clamp_scale="absolute", ) @@ -197,8 +198,15 @@ def _step(self, descriptor: DescrptDPA4) -> tuple[np.ndarray, np.ndarray]: gradient = torch.autograd.grad(output.sum(), coord)[0] return output.detach().cpu().numpy(), gradient.detach().cpu().numpy() - def test_source_gated_flash_retains_dense_rotations(self, monkeypatch) -> None: - """Source-gated flash inference retains the rotations its fallback uses.""" + def test_source_gated_inference_takes_the_fused_convolution( + self, monkeypatch + ) -> None: + """A bridged descriptor runs the fused convolution and matches the dense path. + + The source gate is folded into the edge envelope, so the fused + convolution and the flash attention see it through the one factor they + already weight by, and the dense per-edge rotations are not built. + """ if not so2_conv.op_available(): pytest.skip("the DPA4 CUDA inference operators are unavailable") if not SO2_VALUE_PATH_TRITON_AVAILABLE: @@ -234,8 +242,8 @@ def test_source_gated_flash_retains_dense_rotations(self, monkeypatch) -> None: pytest.skip("the descriptor layout has no fused CUDA convolution") assert conv._flash_atten_fn is not None assert conv._cuda_value_train is None - assert not accelerated._wigner_free_conv - assert accelerated._build_full_wigner() + assert accelerated._wigner_free_conv + assert not accelerated._build_full_wigner() output, gradient = self._step(accelerated) np.testing.assert_allclose(output, dense_output, rtol=2e-4, atol=2e-5) diff --git a/source/tests/pt_expt/descriptor/test_dpa4c.py b/source/tests/pt_expt/descriptor/test_dpa4c.py index 8e59ee2c35..4eef728456 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4c.py +++ b/source/tests/pt_expt/descriptor/test_dpa4c.py @@ -29,12 +29,21 @@ ) # Structural variants exercised by the backend-agnostic contracts. They span -# both angular profiles and both radial function classes. +# both angular profiles and both radial function classes, and one bridging +# window whose transition zone holds several pairs of the test cluster. STRUCTURES = [ {"channels": 16, "lmax": 2, "radial_modes": 0}, {"channels": 32, "lmax": 4, "radial_modes": 0}, {"channels": 32, "lmax": 2, "radial_modes": 3}, {"channels": 16, "lmax": 3, "radial_modes": 3}, + { + "channels": 16, + "lmax": 3, + "radial_modes": 3, + "inner_clamp_f_inner": 1.0, + "inner_clamp_f_outer": 1.25, + "inner_clamp_scale": "absolute", + }, ] @@ -239,7 +248,11 @@ def test_compression_round_trip_preserves_frozen_descriptor(self) -> None: assert restored.compress assert not tuple(restored.radial_embedding.parameters()) assert not any(parameter.requires_grad for parameter in restored.parameters()) - torch.testing.assert_close(self._evaluate(restored, self.coord), reference) + # The descriptor computes in single precision, so the comparison takes + # the single-precision tolerances whatever the dtype of the output. + torch.testing.assert_close( + self._evaluate(restored, self.coord), reference, rtol=1.3e-6, atol=1e-5 + ) def test_serialization_preserves_parameters(self) -> None: restored = DescrptDPA4C.deserialize(self.descriptor.serialize()).to(env.DEVICE) @@ -456,6 +469,56 @@ def test_cutoff_edge_matches_removed_topology(self) -> None: rtol=0.0, ) + def test_bridging_window_is_c3_continuous(self) -> None: + window = { + "inner_clamp_f_inner": 0.5, + "inner_clamp_f_outer": 0.8, + "inner_clamp_scale": "absolute", + } + descriptor = self.build(channels=16, lmax=4, **window).eval() + plain = self.build(channels=16, lmax=4).eval() + # The septic switch meets both radii with three vanishing derivatives: + # at the inner radius the probe joins the constant of the removed + # edge, at the outer radius it joins the plain descriptor. The + # regularized distance exceeds the separation by 1e-14 Å, which the + # fourth derivative of the switch turns into the tolerance below. + inner = self._dimer_derivatives(descriptor, 0.5) + for order in (1, 2, 3): + torch.testing.assert_close( + inner[order], + torch.zeros_like(inner[order]), + atol=1e-8, + rtol=0.0, + ) + outer = self._dimer_derivatives(descriptor, 0.8) + reference = self._dimer_derivatives(plain, 0.8) + assert reference[3].abs() > 1e-3 + for order in range(4): + torch.testing.assert_close( + outer[order], + reference[order], + atol=1e-8, + rtol=0.0, + ) + + def test_frozen_edge_matches_removed_topology(self) -> None: + descriptor = self.build( + inner_clamp_f_inner=0.5, + inner_clamp_f_outer=0.8, + inner_clamp_scale="absolute", + ).eval() + plain = self.build().eval() + radius = torch.tensor(0.5, dtype=torch.float64, device=env.DEVICE) + removed = self._dimer_probe(descriptor, radius, active=False) + # The plain descriptor must see the pair, or the comparison is vacuous. + assert (self._dimer_probe(plain, radius) - removed).abs() > 1e-3 + torch.testing.assert_close( + self._dimer_probe(descriptor, radius, active=True), + removed, + atol=0.0, + rtol=0.0, + ) + def test_coincident_edge_has_finite_third_derivative(self) -> None: descriptor = self.build(lmax=3).eval() derivatives = self._dimer_derivatives(descriptor, 0.0) diff --git a/source/tests/pt_expt/descriptor/test_dpa4c_cpu.py b/source/tests/pt_expt/descriptor/test_dpa4c_cpu.py index d7c7c0c23b..38c885b6f6 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4c_cpu.py +++ b/source/tests/pt_expt/descriptor/test_dpa4c_cpu.py @@ -9,6 +9,10 @@ different compiled body or a different addressing mode. """ +from typing import ( + Any, +) + import pytest import torch @@ -25,6 +29,7 @@ from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( _reference_descriptor, build_compression_artifacts, + contact_radius_input, descriptor_profile, ensure_registered, op_available, @@ -52,10 +57,22 @@ ) +#: Element symbols of the two test types. They fix the covalent radii, whose +#: pairwise sums give the three distinct length scales the window is measured +#: against, so a kernel that gathered the wrong radius would misplace it. +_TYPE_MAP = ["O", "H"] + +#: A bridging window wide enough that the random test cluster populates its +#: frozen zone, its transition zone and the plain zone beyond it. The bounds +#: are fractions of each pair's own length scale, not distances in Å. +_WINDOW = {"inner_clamp_f_inner": 0.9, "inner_clamp_f_outer": 1.6} + + def _build_descriptor( channels: int, lmax: int = 2, radial_modes: int = 0, + **window: float, ) -> DescrptDPA4C: return DescrptDPA4C( rcut=3.0, @@ -66,9 +83,23 @@ def _build_descriptor( radial_modes=radial_modes, precision="float32", seed=17, + type_map=_TYPE_MAP, + **window, ).eval() +def _reduced_distance( + descriptor: DescrptDPA4C, + graph: Any, + atype: torch.Tensor, +) -> torch.Tensor: + """Return the valid edge distances in units of each pair's length scale.""" + source, destination = graph.edge_index[0], graph.edge_index[1] + distance = graph.edge_vec.norm(dim=-1, keepdim=True) + contact = descriptor.pair_contact(atype[destination], atype[source], distance.dtype) + return (distance / contact)[graph.edge_mask, 0] + + def test_fixed_basis_types_freeze_the_basis() -> None: """``/fix`` basis types keep the DPA4C basis parameters out of training.""" for family in ("bessel", "gaussian"): @@ -80,6 +111,7 @@ def test_fixed_basis_types_freeze_the_basis() -> None: n_radial=8, precision="float32", seed=17, + type_map=_TYPE_MAP, basis_type=f"{family}/fix", ) assert fixed.radial_basis.basis_family == family @@ -100,6 +132,7 @@ def test_adam_route_patterns_name_the_first_radial_layer() -> None: radial_modes=2, precision="float32", seed=17, + type_map=_TYPE_MAP, ) patterns = descriptor.adam_route_patterns() assert patterns == ["radial_embedding.layers.0."] @@ -152,12 +185,28 @@ def _arguments(descriptor: DescrptDPA4C, graph, atype: torch.Tensor) -> tuple: artifacts["spin_type"][:0], artifacts["spin_pair"], artifacts["spin_type"], + contact_radius_input(descriptor), bool(graph.destination_sorted), int(descriptor.lmax), *(float(value) for value in artifacts["info"]), + descriptor.bridging_f_inner or 0.0, + descriptor.bridging_f_outer or 0.0, ) +def _descriptor_backward( + cotangent: torch.Tensor, + state: torch.Tensor, + edge_vec: torch.Tensor, + *arguments: Any, +) -> tuple[torch.Tensor, ...]: + """Differentiate the descriptor alone, without an analytical pair potential.""" + no_pair = edge_vec.new_empty(0) + return torch.ops.deepmd.dpa4c_graph_compress_backward( + cotangent, state, edge_vec, *arguments, no_pair, no_pair + )[:3] + + def _spin_free(arguments: tuple) -> tuple: """Drop the native spin block, which the CPU kernel does not implement.""" return (*arguments[:15], *arguments[18:]) @@ -181,9 +230,7 @@ def _check_parity(descriptor: DescrptDPA4C, canonical: bool) -> None: (reference_gradient,) = torch.autograd.grad( (reference * cotangent).sum(), reference_edge ) - gradient = torch.ops.deepmd.dpa4c_graph_compress_backward( - cotangent, state, graph.edge_vec, *arguments - )[0] + gradient = _descriptor_backward(cotangent, state, graph.edge_vec, *arguments)[0] torch.testing.assert_close(gradient, reference_gradient, atol=8.0e-6, rtol=1.0e-4) @@ -195,6 +242,97 @@ def test_forward_backward_parity(channels: int, canonical: bool) -> None: _check_parity(_build_descriptor(channels), canonical) +@_CPU +@pytest.mark.parametrize("canonical", [False, True]) +def test_bridging_window_parity(canonical: bool) -> None: + """The switched envelope reproduces the portable descriptor. + + The reference evaluates the window in its own form, so it is first pinned + to the descriptor equations and then serves the kernel comparison. + """ + descriptor = _build_descriptor(32, lmax=3, radial_modes=2, **_WINDOW) + graph, atype = _build_graph(descriptor, canonical) + reduced = _reduced_distance(descriptor, graph, atype) + inner, outer = _WINDOW.values() + assert (reduced < inner).any() and (reduced > outer).any() + assert ((reduced > inner) & (reduced < outer)).sum() > 8 + reference = _reference_descriptor( + graph.edge_vec, *_spin_free(_arguments(descriptor, graph, atype)) + ) + portable, _ = descriptor.call_graph( + graph, atype, type_embedding=descriptor.type_embedding.call() + ) + torch.testing.assert_close(reference, portable, atol=3.0e-5, rtol=3.0e-5) + _check_parity(descriptor, canonical) + + +@_CPU +@pytest.mark.parametrize("mode", ["zbl", "nlh"]) +@pytest.mark.parametrize("windowed", [True, False]) +@pytest.mark.parametrize("canonical", [False, True]) +def test_pair_potential_matches_the_portable_term( + canonical: bool, + windowed: bool, + mode: str, +) -> None: + """The backward scan evaluates the analytical pair potential of a bridged model. + + A vanishing descriptor cotangent isolates the term: the edge gradient is + then the pair slope alone, weighted by the energy cotangent of the + destination, and the scan returns the pair energy of every node. Edges + inside the bridging window leave the descriptor but keep the pair term, + and the portable reference backs the traced form of the fused operator. + Without a window the pair table alone selects the bridged backward, behind + a plain forward. The scan reads the coefficients of whichever potential + filled the table, so both supported modes ride the same kernel. + """ + from deepmd.dpmodel.atomic_model.inner_potential import ( + InnerPotential, + ) + from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( + _reference_backward, + ) + + descriptor = _build_descriptor(32, **(_WINDOW if windowed else {})) + graph, atype = _build_graph(descriptor, canonical) + arguments = _arguments(descriptor, graph, atype) + output, state = torch.ops.deepmd.dpa4c_graph_compress(graph.edge_vec, *arguments) + potential = InnerPotential(type_map=["O", "H"], mode=mode) + pair_table = torch.as_tensor(potential.pair_table) + # Ghost nodes carry no energy, so their cotangent vanishes. + seed = (torch.arange(atype.shape[0], device="cpu") % 3 != 0).to(torch.float64) + + edge = graph.edge_vec.double().requires_grad_(True) + reference = potential.call( + edge, graph.edge_index, atype, graph.edge_mask, atype.shape[0] + )[0] + (reference_gradient,) = torch.autograd.grad((reference[:, 0] * seed).sum(), edge) + assert reference.max() > 10.0 + for backward in ( + torch.ops.deepmd.dpa4c_graph_compress_backward, + _reference_backward, + ): + edge_gradient, _, _, pair_energy = backward( + torch.zeros_like(output), + state, + graph.edge_vec, + *arguments, + pair_table, + seed, + ) + torch.testing.assert_close(pair_energy, reference, atol=3e-5, rtol=3e-5) + torch.testing.assert_close( + edge_gradient.double(), reference_gradient, atol=3e-5, rtol=3e-5 + ) + + no_pair = graph.edge_vec.new_empty(0) + plain_gradient, _, _, absent = torch.ops.deepmd.dpa4c_graph_compress_backward( + torch.zeros_like(output), state, graph.edge_vec, *arguments, no_pair, no_pair + ) + assert absent.numel() == 0 + assert torch.count_nonzero(plain_gradient).item() == 0 + + @_CPU @pytest.mark.parametrize("lmax", [2, 3, 4]) @pytest.mark.parametrize("radial_modes", [0, 2, 4, 8]) @@ -227,9 +365,7 @@ def test_masked_edges_are_ignored() -> None: torch.testing.assert_close(output, reference) cotangent = torch.ones_like(output) - gradient = torch.ops.deepmd.dpa4c_graph_compress_backward( - cotangent, state, padded, *extended - )[0] + gradient = _descriptor_backward(cotangent, state, padded, *extended)[0] assert torch.all(gradient[graph.edge_vec.shape[0] :] == 0.0) @@ -496,6 +632,7 @@ def test_spin_descriptor_keeps_the_reference_path() -> None: n_radial=8, precision="float32", seed=17, + type_map=_TYPE_MAP, use_spin=[True, False], ).eval() assert mega_eligible(with_spin) diff --git a/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py index 5734b011fa..d1586c5856 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py +++ b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py @@ -5,6 +5,9 @@ from collections.abc import ( Sequence, ) +from typing import ( + Any, +) import numpy as np import pytest @@ -27,6 +30,7 @@ _table_lookup, build_compression_artifacts, build_radial_table, + contact_radius_input, descriptor_profile, dpa4c_graph_compress_energy_force, ensure_registered, @@ -40,10 +44,22 @@ ) +#: Element symbols of the two test types. They fix the covalent radii, whose +#: pairwise sums give the three distinct length scales the window is measured +#: against, so a kernel that gathered the wrong radius would misplace it. +_TYPE_MAP = ["O", "H"] + +#: A bridging window wide enough that the random test cluster populates its +#: frozen zone, its transition zone and the plain zone beyond it. The bounds +#: are fractions of each pair's own length scale, not distances in Å. +_WINDOW = {"inner_clamp_f_inner": 0.9, "inner_clamp_f_outer": 1.6} + + def _build_descriptor( channels: int, lmax: int = 2, radial_modes: int = 0, + **window: float, ) -> DescrptDPA4C: return ( DescrptDPA4C( @@ -55,12 +71,26 @@ def _build_descriptor( radial_modes=radial_modes, precision="float32", seed=17, + type_map=_TYPE_MAP, + **window, ) .cuda() .eval() ) +def _reduced_distance( + descriptor: DescrptDPA4C, + graph: Any, + atype: torch.Tensor, +) -> torch.Tensor: + """Return the valid edge distances in units of each pair's length scale.""" + source, destination = graph.edge_index[0], graph.edge_index[1] + distance = graph.edge_vec.norm(dim=-1, keepdim=True) + contact = descriptor.pair_contact(atype[destination], atype[source], distance.dtype) + return (distance / contact)[graph.edge_mask, 0] + + def _build_graph( descriptor: DescrptDPA4C, canonical: bool, @@ -125,12 +155,28 @@ def _arguments( artifacts["spin_type"][:0], artifacts["spin_pair"], artifacts["spin_type"], + contact_radius_input(descriptor), bool(graph.destination_sorted), int(descriptor.lmax), *(float(value) for value in artifacts["info"]), + descriptor.bridging_f_inner or 0.0, + descriptor.bridging_f_outer or 0.0, ) +def _descriptor_backward( + cotangent: torch.Tensor, + state: torch.Tensor, + edge_vec: torch.Tensor, + *arguments: Any, +) -> tuple[torch.Tensor, ...]: + """Differentiate the descriptor alone, without an analytical pair potential.""" + no_pair = edge_vec.new_empty(0) + return torch.ops.deepmd.dpa4c_graph_compress_backward( + cotangent, state, edge_vec, *arguments, no_pair, no_pair + )[:3] + + def _spin_free(arguments: tuple) -> tuple: """Drop the native spin block for the spin-free CPU reference.""" return (*arguments[:15], *arguments[18:]) @@ -250,13 +296,13 @@ def test_wide_backward_is_deterministic(channels: int) -> None: previous = torch.are_deterministic_algorithms_enabled() torch.use_deterministic_algorithms(True) try: - first = torch.ops.deepmd.dpa4c_graph_compress_backward( + first = _descriptor_backward( torch.ones_like(output), state, graph.edge_vec, *arguments, ) - second = torch.ops.deepmd.dpa4c_graph_compress_backward( + second = _descriptor_backward( torch.ones_like(output), state, graph.edge_vec, @@ -268,10 +314,19 @@ def test_wide_backward_is_deterministic(channels: int) -> None: @_GPU +@pytest.mark.parametrize("window", [{}, _WINDOW], ids=["plain", "bridged"]) @pytest.mark.parametrize("channels", [8, 16, 32, 64, 128]) -def test_compressed_matches_uncompressed_descriptor(channels: int) -> None: - descriptor = _build_descriptor(channels) +def test_compressed_matches_uncompressed_descriptor( + channels: int, + window: dict[str, float], +) -> None: + descriptor = _build_descriptor(channels, **window) graph, atype = _build_graph(descriptor, canonical=False) + if window: + reduced = _reduced_distance(descriptor, graph, atype) + inner, outer = window.values() + assert (reduced < inner).any() and (reduced > outer).any() + assert ((reduced > inner) & (reduced < outer)).sum() > 8 arguments = _arguments(descriptor, graph, atype) compressed, _state = torch.ops.deepmd.dpa4c_graph_compress( graph.edge_vec, @@ -431,6 +486,7 @@ def _build_charge_descriptor( radial_modes=radial_modes, precision="float32", seed=17, + type_map=_TYPE_MAP, add_chg_spin_ebd=True, default_chg_spin=( None if default_chg_spin is None else list(default_chg_spin) @@ -544,6 +600,7 @@ def test_compression_rejects_float64_descriptor() -> None: n_radial=8, precision="float64", seed=17, + type_map=_TYPE_MAP, ) .cuda() .eval() @@ -557,19 +614,24 @@ def test_compression_rejects_float64_descriptor() -> None: @pytest.mark.parametrize("lmax", [2, 3, 4]) @pytest.mark.parametrize("radial_modes", [0, 2, 4, 8]) @pytest.mark.parametrize("canonical", [False, True]) +@pytest.mark.parametrize("windowed", [False, True]) def test_supported_surface_parity( channels: int, lmax: int, radial_modes: int, canonical: bool, + windowed: bool, ) -> None: """Cover the complete compiled surface against the portable equations. Each scalar width owns a distinct lane mapping and each angular degree a distinct instantiation, so the cross product is the contract the operator - advertises rather than a sample of it. + advertises rather than a sample of it. A bridging window selects the + bridged specialization of every profile. """ - descriptor = _build_descriptor(channels, lmax, radial_modes) + descriptor = _build_descriptor( + channels, lmax, radial_modes, **(_WINDOW if windowed else {}) + ) graph, atype = _build_graph(descriptor, canonical) arguments = _arguments(descriptor, graph, atype) ensure_registered() @@ -593,7 +655,10 @@ def test_supported_surface_parity( (reference_value * cotangent).sum(), reference_edge, ) - torch.testing.assert_close(gradient, reference_gradient, atol=8e-6, rtol=1e-4) + # Inside the window the switch is steep, and the float32 reference then + # carries a rounding error of its own that exceeds the plain tolerance. + atol, rtol = (1e-4, 5e-4) if windowed else (8e-6, 1e-4) + torch.testing.assert_close(gradient, reference_gradient, atol=atol, rtol=rtol) @_GPU @@ -778,16 +843,25 @@ def test_compressed_cutoff_matches_removed_topology(channels: int) -> None: @_GPU +@pytest.mark.parametrize("frozen_by", ["mask", "window"]) @pytest.mark.parametrize("channels", [8, 64, 128]) -def test_in_row_mask_matches_removed_edge(channels: int) -> None: - descriptor = _build_descriptor(channels) +def test_in_row_mask_matches_removed_edge(channels: int, frozen_by: str) -> None: + """An edge switched off in its row equals an edge absent from the graph. + + The mask and the bridging window are the two ways a row may hold an edge + that contributes nothing; the second edge lies inside the frozen zone of + the window and is switched off by either. + """ + descriptor = _build_descriptor( + channels, **(_WINDOW if frozen_by == "window" else {}) + ) edge_index = torch.tensor( [[1, 0], [0, 1]], dtype=torch.long, device="cuda", ) edge_vec = torch.tensor( - [[1.0, 0.2, -0.1], [-0.7, 0.3, 0.4]], + [[1.0, 1.4, -0.1], [-0.7, 0.3, 0.4]], dtype=torch.float32, device="cuda", ) @@ -806,7 +880,9 @@ def test_in_row_mask_matches_removed_edge(channels: int) -> None: ) masked_graph = dataclasses.replace( graph, - edge_mask=torch.tensor([True, False], dtype=torch.bool, device="cuda"), + edge_mask=torch.tensor( + [True, frozen_by == "window"], dtype=torch.bool, device="cuda" + ), ) masked_edge = edge_vec.detach().clone().requires_grad_(True) masked, _masked_state = torch.ops.deepmd.dpa4c_graph_compress( @@ -842,7 +918,8 @@ def test_in_row_mask_matches_removed_edge(channels: int) -> None: removed_edge, ) - assert not torch.allclose(full, masked) + if frozen_by == "mask": + assert not torch.allclose(full, masked) torch.testing.assert_close(masked, removed, atol=2e-6, rtol=2e-6) torch.testing.assert_close( masked_gradient[:1], @@ -859,10 +936,101 @@ def test_in_row_mask_matches_removed_edge(channels: int) -> None: assert torch.count_nonzero(masked_gradient[0]).item() > 0 +@_GPU +@pytest.mark.parametrize("mode", ["zbl", "nlh"]) +@pytest.mark.parametrize("windowed", [True, False]) +@pytest.mark.parametrize("canonical", [False, True]) +@pytest.mark.parametrize("channels", [8, 32, 128]) +@pytest.mark.parametrize(("lmax", "radial_modes"), [(2, 0), (3, 4), (4, 2)]) +def test_pair_potential_matches_the_portable_term( + channels: int, + canonical: bool, + windowed: bool, + lmax: int, + radial_modes: int, + mode: str, +) -> None: + """The backward scan evaluates the analytical pair potential of a bridged model. + + A vanishing descriptor cotangent isolates the term: the edge gradient is + then the pair slope alone, weighted by the energy cotangent of the + destination, and the scan returns the pair energy of every node. The three + widths cover the two-, four- and eight-lane splits of the four series + terms. Without a window the pair table alone selects the bridged backward, + behind a plain forward. Both supported modes ride the kernel, which reads + the coefficients of whichever potential filled the table; their decay rates + differ by three orders of magnitude, so the single-precision table has to + carry either. + """ + from deepmd.dpmodel.atomic_model.inner_potential import ( + InnerPotential, + ) + + descriptor = _build_descriptor( + channels, lmax, radial_modes, **(_WINDOW if windowed else {}) + ) + graph, atype = _build_graph(descriptor, canonical=canonical) + arguments = _arguments(descriptor, graph, atype) + output, state = torch.ops.deepmd.dpa4c_graph_compress(graph.edge_vec, *arguments) + potential = InnerPotential(type_map=["O", "H"], mode=mode) + pair_table = torch.as_tensor(potential.pair_table, device="cuda") + # Ghost nodes carry no energy, so their cotangent vanishes. + seed = (torch.arange(atype.shape[0], device="cuda") % 3 != 0).to(torch.float64) + + edge_gradient, _, _, pair_energy = torch.ops.deepmd.dpa4c_graph_compress_backward( + torch.zeros_like(output), + state, + graph.edge_vec, + *arguments, + pair_table, + seed, + ) + edge = graph.edge_vec.double().requires_grad_(True) + reference = potential.call( + edge, graph.edge_index, atype, graph.edge_mask, atype.shape[0] + )[0] + (reference_gradient,) = torch.autograd.grad((reference[:, 0] * seed).sum(), edge) + assert reference.max() > 10.0 + torch.testing.assert_close(pair_energy, reference, atol=3e-5, rtol=3e-5) + torch.testing.assert_close( + edge_gradient.double(), reference_gradient, atol=3e-5, rtol=3e-5 + ) + + no_pair = graph.edge_vec.new_empty(0) + plain_gradient, _, _, absent = torch.ops.deepmd.dpa4c_graph_compress_backward( + torch.zeros_like(output), state, graph.edge_vec, *arguments, no_pair, no_pair + ) + assert absent.numel() == 0 + assert torch.count_nonzero(plain_gradient).item() == 0 + + # A pair term without energy cotangent leaves the descriptor gradient of + # the plain backward. Without a window that gradient comes from the plain + # specialization, which the compiler schedules and contracts independently + # of the bridged one, so the two agree to single-precision rounding, far + # below any slope of the pair term. + cotangent = torch.randn_like(output) + descriptor_gradient = _descriptor_backward( + cotangent, state, graph.edge_vec, *arguments + )[0] + unseeded_gradient = torch.ops.deepmd.dpa4c_graph_compress_backward( + cotangent, + state, + graph.edge_vec, + *arguments, + pair_table, + torch.zeros_like(seed), + )[0] + scale = descriptor_gradient.abs().max().item() + torch.testing.assert_close( + unseeded_gradient, descriptor_gradient, atol=1e-6 * scale, rtol=1e-6 + ) + + @_GPU @pytest.mark.parametrize("channels", [8, 32]) -def test_padding_type_edge_has_zero_gradient(channels: int) -> None: - descriptor = _build_descriptor(channels) +@pytest.mark.parametrize("bridged", [False, True]) +def test_padding_type_edge_has_zero_gradient(channels: int, bridged: bool) -> None: + descriptor = _build_descriptor(channels, **(_WINDOW if bridged else {})) edge_vec = torch.tensor( [[1.0, 0.2, -0.1]], dtype=torch.float32, @@ -912,17 +1080,19 @@ def test_int32_edge_indices() -> None: @_GPU +@pytest.mark.parametrize("windowed", [False, True]) @pytest.mark.parametrize("channels", [8, 64, 128]) @pytest.mark.parametrize("index_dtype", [torch.int64, torch.uint32]) def test_compact_canonical_parity( channels: int, index_dtype: torch.dtype, + windowed: bool, ) -> None: from deepmd.pt_expt.kernels.dpa4c.canonical import ( ensure_registered as ensure_canonical_registered, ) - descriptor = _build_descriptor(channels) + descriptor = _build_descriptor(channels, **(_WINDOW if windowed else {})) graph, atype = _build_graph(descriptor, canonical=True) arguments = _arguments(descriptor, graph, atype) ensure_canonical_registered() @@ -934,8 +1104,8 @@ def test_compact_canonical_parity( graph.edge_index[0].to(index_dtype), graph.destination_row_ptr, atype, - *arguments[5:18], - *arguments[19:], + *arguments[5:19], + *arguments[20:], ) compact_output, compact_state = torch.ops.deepmd.dpa4c_canonical_compress( graph.edge_vec, @@ -945,7 +1115,7 @@ def test_compact_canonical_parity( torch.testing.assert_close(compact_state, generic_state) cotangent = torch.randn_like(generic_output) - generic_gradient = torch.ops.deepmd.dpa4c_graph_compress_backward( + generic_gradient = _descriptor_backward( cotangent, generic_state, graph.edge_vec, @@ -979,8 +1149,8 @@ def test_compact_inplace_backward_reuses_state(channels: int) -> None: graph.edge_index[0].to(torch.uint32), graph.destination_row_ptr, atype, - *arguments[5:18], - *arguments[19:], + *arguments[5:19], + *arguments[20:], ) ensure_canonical_registered() output, state = torch.ops.deepmd.dpa4c_canonical_compress( @@ -1214,11 +1384,105 @@ def run() -> tuple[torch.Tensor, ...]: torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-6) +@_GPU +@pytest.mark.parametrize("bridged", [False, True]) +def test_fused_canonical_cpu_reference_matches_the_kernel( + bridged: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The CPU implementation of the fused canonical operator reproduces it. + + The operator runs once on CUDA through the energy-force helper and once on + CPU copies of the same arguments, where it dispatches to the reference + sequence of descriptor, fitting, fitting backward and descriptor backward, + with the pair potential when the model is bridged. + """ + from deepmd.dpmodel.atomic_model.inner_potential import ( + InnerPotential, + ) + from deepmd.pt_expt.fitting.ener_fitting import ( + EnergyFittingNet, + ) + from deepmd.pt_expt.kernels.dpa4c import ( + canonical, + ) + from deepmd.pt_expt.kernels.dpa4c.canonical import ( + dpa4c_canonical_compress_energy_force, + ) + from deepmd.pt_expt.utils.canonical_graph import ( + canonical_graph_from_neighbor_graph, + ) + + canonical.ensure_registered() + window = {"inner_clamp_f_inner": 0.9, "inner_clamp_f_outer": 1.3} if bridged else {} + descriptor = _build_descriptor(8, **window) + neighbor_graph, atype = _build_graph(descriptor, canonical=True) + graph = canonical_graph_from_neighbor_graph( + dataclasses.replace(neighbor_graph, n_local=neighbor_graph.n_node) + ) + descriptor._set_compression(build_compression_artifacts(descriptor)) + fitting = ( + EnergyFittingNet( + ntypes=2, + dim_descrpt=descriptor.get_dim_out(), + neuron=[32, 32], + resnet_dt=False, + activation_function="silu", + precision="float32", + mixed_types=True, + seed=29, + ) + .cuda() + .eval() + ) + pair_table = ( + torch.as_tensor(InnerPotential(["O", "H"]).pair_table, device="cuda") + if bridged + else None + ) + operator = torch.ops.deepmd.dpa4c_canonical_compress_energy_gradient + calls = [] + + def record(*args: Any) -> tuple[torch.Tensor, ...]: + calls.append(args) + return operator(*args) + + monkeypatch.setattr( + torch.ops.deepmd, "dpa4c_canonical_compress_energy_gradient", record + ) + dpa4c_canonical_compress_energy_force( + descriptor, + fitting, + graph, + atype, + torch.ones(atype.shape[0], dtype=torch.bool, device="cuda"), + fitting.bias_atom_e[:, 0], + False, + None, + pair_table, + ) + (arguments,) = calls + + def to_cpu(value: Any) -> Any: + if isinstance(value, torch.Tensor): + return value.cpu() + if isinstance(value, list): + return [to_cpu(item) for item in value] + return value + + kernel = operator(*arguments) + reference = operator(*[to_cpu(argument) for argument in arguments]) + for actual, expected in zip(reference, kernel, strict=True): + scale = max(expected.abs().max().item(), 1.0) if expected.numel() else 1.0 + torch.testing.assert_close(actual, expected.cpu(), atol=2e-5 * scale, rtol=2e-5) + + def _build_spin_descriptor( channels: int, lmax: int = 2, radial_modes: int = 0, device: str = "cuda", + **window: float, ) -> DescrptDPA4C: """Return a spin-conditioned descriptor with a non-unit reference moment. @@ -1238,7 +1502,9 @@ def _build_spin_descriptor( radial_modes=radial_modes, precision="float32", seed=17, + type_map=_TYPE_MAP, use_spin=[True, False], + **window, ) .to(device) .eval() @@ -1302,13 +1568,11 @@ def test_empty_native_spin_cuda_backward_preserves_spin_contract() -> None: *arguments, ) - edge_gradient, spin_gradient, edge_spin_gradient = ( - torch.ops.deepmd.dpa4c_graph_compress_backward( - torch.empty_like(output), - state, - graph.edge_vec, - *arguments, - ) + edge_gradient, spin_gradient, edge_spin_gradient = _descriptor_backward( + torch.empty_like(output), + state, + graph.edge_vec, + *arguments, ) assert edge_gradient.shape == (0, 3) @@ -1321,11 +1585,13 @@ def test_empty_native_spin_cuda_backward_preserves_spin_contract() -> None: ("channels", "lmax", "radial_modes"), [(8, 2, 0), (16, 2, 0), (32, 2, 4), (64, 3, 0), (128, 3, 8), (32, 4, 0)], ) +@pytest.mark.parametrize("windowed", [False, True]) def test_spin_compressed_matches_portable( monkeypatch: pytest.MonkeyPatch, channels: int, lmax: int, radial_modes: int, + windowed: bool, ) -> None: """Descriptor, coordinate gradient and magnetic force all match. @@ -1336,7 +1602,9 @@ def test_spin_compressed_matches_portable( their index inflates the gradient magnitude at the widest profile until fp32 tabulation noise alone exceeds the tolerance. """ - descriptor = _build_spin_descriptor(channels, lmax, radial_modes) + descriptor = _build_spin_descriptor( + channels, lmax, radial_modes, **(_WINDOW if windowed else {}) + ) assert mega_eligible(descriptor) graph, atype = _build_graph(descriptor, canonical=True) generator = torch.Generator(device="cuda").manual_seed(11) @@ -1376,8 +1644,11 @@ def run() -> tuple[torch.Tensor, ...]: actual = run() _assert_dispatched(actual[0], reference[0]) # Tabulation error reaches the gradients through the table derivative, so - # they carry the wider tolerance the geometric backward already uses. - tolerances = ((3e-5, 3e-5), (8e-6, 1e-4), (8e-6, 1e-4)) + # they carry the wider tolerance the geometric backward already uses, and + # the steep switch inside a bridging window widens it as in the surface + # test. + gradient = (1e-4, 5e-4) if windowed else (8e-6, 1e-4) + tolerances = ((3e-5, 3e-5), gradient, gradient) for (atol, rtol), value, expected in zip( tolerances, actual, reference, strict=True ): @@ -1494,19 +1765,35 @@ def run(edge_vec: torch.Tensor, moment: torch.Tensor) -> torch.Tensor: @_GPU +@pytest.mark.parametrize("bridged", [False, True]) @pytest.mark.parametrize("spin_conditioned", [False, True]) -def test_backward_operator_satisfies_its_schema(spin_conditioned: bool) -> None: - """The backward operator declares three independent results. +def test_backward_operator_satisfies_its_schema( + spin_conditioned: bool, + bridged: bool, +) -> None: + """The backward operator declares four independent results. - All three are unannotated, so any two of them sharing storage would be an + All four are unannotated, so any two of them sharing storage would be an alias the schema does not describe, which is undefined under functionalization. ``opcheck`` decides that mechanically, including on the - spin-free path where two of the three are absent and an allocation shared - between them would otherwise go unnoticed. + spin-free path without a pair potential, where three of the four are + absent and an allocation shared between them would otherwise go unnoticed. """ + from deepmd.dpmodel.atomic_model.inner_potential import ( + InnerPotential, + ) + descriptor = _build_spin_descriptor(8) if spin_conditioned else _build_descriptor(8) graph, atype = _build_graph(descriptor, canonical=True) arguments = _arguments(descriptor, graph, atype) + pair = ( + ( + torch.as_tensor(InnerPotential(["O", "H"]).pair_table, device="cuda"), + torch.ones(atype.shape[0], dtype=torch.float64, device="cuda"), + ) + if bridged + else (graph.edge_vec.new_empty(0),) * 2 + ) if spin_conditioned: arguments = _with_spin( arguments, @@ -1518,7 +1805,7 @@ def test_backward_operator_satisfies_its_schema(spin_conditioned: bool) -> None: ) torch.library.opcheck( torch.ops.deepmd.dpa4c_graph_compress_backward.default, - (torch.ones_like(output), state, graph.edge_vec, *arguments), + (torch.ones_like(output), state, graph.edge_vec, *arguments, *pair), ) diff --git a/source/tests/pt_expt/fitting/test_ener_fitting.py b/source/tests/pt_expt/fitting/test_ener_fitting.py index fe61daf6bc..02c096fbdb 100644 --- a/source/tests/pt_expt/fitting/test_ener_fitting.py +++ b/source/tests/pt_expt/fitting/test_ener_fitting.py @@ -1,7 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import unittest +from unittest import ( + mock, +) import numpy as np +import pytest import torch from torch.fx.experimental.proxy_tensor import ( make_fx, @@ -10,8 +14,16 @@ from deepmd.dpmodel.descriptor import ( DescrptSeA, ) +from deepmd.dpmodel.fitting.ener_fitting import EnergyFittingNet as EnergyFittingNetDP +from deepmd.pt.cxx_op import ( + ENABLE_CUSTOMIZED_OP, +) from deepmd.pt_expt.fitting import ( EnergyFittingNet, + ener_fitting, +) +from deepmd.pt_expt.kernels.graph_fitting import ( + op_available, ) from deepmd.pt_expt.utils import ( env, @@ -28,6 +40,83 @@ ) +@pytest.mark.skipif( + not ENABLE_CUSTOMIZED_OP or not op_available(), + reason="the fused fitting operator is unavailable", +) +@pytest.mark.parametrize("activation", ["tanh", "silu"]) # activation_function +@pytest.mark.parametrize( + "reference_mode", ["none", "vacuum", "folded"] +) # vacuum reference: none / passed per call / folded into the bias +def test_fused_node_gate_matches_portable( + monkeypatch: pytest.MonkeyPatch, activation: str, reference_mode: str +) -> None: + """The fused readout preserves gates and their gradients across vacuum folding.""" + monkeypatch.setenv("DP_CUDA_INFER", "1") + fitting = ( + EnergyFittingNet( + ntypes=2, + dim_descrpt=12, + neuron=[16, 16], + resnet_dt=False, + activation_function=activation, + precision="float32", + mixed_types=True, + vacuum_ref=reference_mode != "none", + seed=3, + ) + .to(env.DEVICE) + .eval() + ) + generator = torch.Generator(device=env.DEVICE).manual_seed(5) + descriptor = torch.randn( + 8, 12, generator=generator, dtype=torch.float32, device=env.DEVICE + ).requires_grad_(True) + atype = torch.arange(8, dtype=torch.int64, device=env.DEVICE) % 2 + gate = torch.tensor( + [0.0, 1.0, 0.2, 0.65, 0.9, 0.0, 0.4, 1.0], + dtype=torch.float64, + device=env.DEVICE, + requires_grad=True, + ) + vacuum = ( + torch.randn(2, 12, generator=generator, dtype=torch.float32, device=env.DEVICE) + if reference_mode != "none" + else None + ) + with torch.no_grad(): + fitting.bias_atom_e[:, 0].copy_(fitting.bias_atom_e.new_tensor([1.25, -0.75])) + if reference_mode == "folded": + fitting.fold_vacuum_reference(vacuum) + vacuum = None + + kwargs = {"vacuum_descriptor": vacuum, "node_gate": gate} + reference = EnergyFittingNetDP.call_graph(fitting, descriptor, atype, **kwargs)[ + "energy" + ] + with mock.patch.object( + ener_fitting, "graph_fitting", wraps=ener_fitting.graph_fitting + ) as fused: + actual = fitting.call_graph(descriptor, atype, **kwargs)["energy"] + fused.assert_called_once() + torch.testing.assert_close(actual, reference, atol=2e-5, rtol=2e-5) + closed = gate == 0.0 + bias = fitting.readout_reference().to(actual.dtype)[atype] + torch.testing.assert_close(actual[closed], bias[closed], atol=0.0, rtol=0.0) + + cotangent = torch.linspace(-1.0, 1.0, 8, device=env.DEVICE)[:, None] + expected_gradients = torch.autograd.grad( + (reference * cotangent).sum(), (descriptor, gate) + ) + actual_gradients = torch.autograd.grad( + (actual * cotangent).sum(), (descriptor, gate) + ) + torch.testing.assert_close( + actual_gradients, expected_gradients, atol=2e-5, rtol=2e-5 + ) + assert actual_gradients[1].abs().max() > 1e-6 + + class TestEnergyFittingNet(unittest.TestCase, TestCaseSingleFrameWithNlist): def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) diff --git a/source/tests/pt_expt/model/test_dpa4_interop.py b/source/tests/pt_expt/model/test_dpa4_interop.py index d972e3020c..145eaf3bd6 100644 --- a/source/tests/pt_expt/model/test_dpa4_interop.py +++ b/source/tests/pt_expt/model/test_dpa4_interop.py @@ -225,6 +225,39 @@ def test_adam_routing( assert "momentum_buffer" not in optimizer.state[parameters[name]] assert "momentum_buffer" in optimizer.state[parameters[control]] + @pytest.mark.parametrize("child_protection", [None, 0.05]) + def test_pt_spin_hybrid_routes_explicit_env_protection( + self, + child_protection: float | None, + ) -> None: + """The PT factory preserves explicit hybrid-level protection.""" + descriptor = { + "type": "se_e2_a", + "sel": [4, 4], + "rcut": 6.0, + "rcut_smth": 5.5, + "neuron": [4, 8], + "axis_neuron": 2, + } + if child_protection is not None: + descriptor["env_protection"] = child_protection + config = { + "type": "standard", + "type_map": ["Fe"], + "descriptor": { + "type": "hybrid", + "env_protection": 0.03, + "list": [descriptor], + }, + "fitting_net": {"type": "ener", "neuron": [8]}, + "spin": {"use_spin": [True], "virtual_scale": [0.2]}, + } + original = copy.deepcopy(config) + model = pt_get_model(config) + expected = 0.03 if child_protection is None else child_protection + assert model.get_descriptor().get_env_protection() == expected + assert config == original + @pytest.mark.parametrize("basis_type", ["bessel", "gaussian"]) def test_single_envelope_normalization_and_roundtrip(self, basis_type: str) -> None: """Preserve the integer envelope configuration and its energy/force function.""" diff --git a/source/tests/pt_expt/model/test_dpa4_native_spin.py b/source/tests/pt_expt/model/test_dpa4_native_spin.py index 4afdae6d47..2510855e8f 100644 --- a/source/tests/pt_expt/model/test_dpa4_native_spin.py +++ b/source/tests/pt_expt/model/test_dpa4_native_spin.py @@ -1213,8 +1213,8 @@ def test_legacy_state_squares_the_env_seed_spin_gate(self) -> None: target = self._descriptor() target.load_state_dict(state) assert torch.all(target.env_seed_embedding.spin_scale == 4.0) - assert target.version == 1.2 - assert float(target.version_tensor.item()) == 1.2 + assert target.version == 1.3 + assert float(target.version_tensor.item()) == 1.3 def test_legacy_spin_free_state_zeros_dormant_routes(self) -> None: state = self._descriptor([False, False]).state_dict() @@ -1238,7 +1238,7 @@ def test_legacy_spin_free_state_zeros_dormant_routes(self) -> None: for key in dormant_keys: assert torch.all(migrated[key] == 0.0) assert torch.all(migrated["spin_embedding.mag_layer1.w"] == 5.0) - assert target.version == 1.2 + assert target.version == 1.3 def test_state_without_a_version_is_read_as_the_last_untagged_one(self) -> None: """Checkpoints predating the buffer were written under version 1.1.""" @@ -1250,7 +1250,7 @@ def test_state_without_a_version_is_read_as_the_last_untagged_one(self) -> None: target = self._descriptor() target.load_state_dict(state) assert torch.all(target.env_seed_embedding.spin_scale == 9.0) - assert target.version == 1.2 + assert target.version == 1.3 class TestNativeSpinConfigFormsPtExpt: diff --git a/source/tests/pt_expt/model/test_dpa4_vacuum_ref.py b/source/tests/pt_expt/model/test_dpa4_vacuum_ref.py index 2dc3147fe7..3a4ee60793 100644 --- a/source/tests/pt_expt/model/test_dpa4_vacuum_ref.py +++ b/source/tests/pt_expt/model/test_dpa4_vacuum_ref.py @@ -228,7 +228,7 @@ def test_dense_and_graph_routes_share_the_vacuum_descriptor() -> None: graph, atype_all, charge_spin, spin = am.append_vacuum_frames( graph, atype.reshape(-1), charge_spin, spin ) - gg, _ = am.descriptor.call_graph( + gg, _, _ = am.descriptor.call_graph( graph, atype_all, type_embedding=am.descriptor.graph_type_embedding_table(), diff --git a/source/tests/pt_expt/model/test_dpa4c_graph_lower.py b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py index 2386089db9..02af9d412d 100644 --- a/source/tests/pt_expt/model/test_dpa4c_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py @@ -377,6 +377,54 @@ def test_shared_weights_preserve_each_tasks_vacuum_reference() -> None: torch.testing.assert_close(actual[name], expected[name]) +def test_vacuum_descriptor_never_materializes_the_neighbor_capacity( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The isolated atoms are evaluated on the graph route, without padding. + + DPA4C reports an unbounded neighbor capacity, so a dense neighbor list of + that width would hold one padded edge per slot and per type; the graph + route holds no real edge. The result must equal the reference rows the + training route appends to a real graph. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, + ) + + config = _config() + config["type_map"] = ["O", "H"] + config["preset_out_bias"] = {"energy": {"O": -3.0, "H": 0.5}} + config["descriptor"].update({"channels": 8, "lmax": 2}) + config["fitting_net"].update({"neuron": [8, 8], "vacuum_ref": True}) + model = get_model(config).to(env.DEVICE).eval() + am = model.atomic_model + assert am.fitting_net.needs_vacuum_descriptor() + + def dense_call(*args: object, **kwargs: object) -> None: + raise AssertionError("the dense adapter must not serve the vacuum descriptor") + + monkeypatch.setattr(am.descriptor, "call", dense_call) + vacuum = am.vacuum_descriptor() + assert vacuum.shape == (2, am.descriptor.get_dim_out()) + + coord = torch.tensor( + [[[0.0, 0.0, 0.0], [1.1, 0.2, -0.1], [-0.4, 0.9, 0.3]]], + dtype=torch.float64, + device=env.DEVICE, + ) + atype = torch.tensor([[0, 1, 1]], dtype=torch.long, device=env.DEVICE) + graph = build_neighbor_graph( + coord, atype, None, config["descriptor"]["rcut"], with_csr=True + ) + graph, atype_all, _, _ = am.append_vacuum_frames( + graph, atype.reshape(-1), None, None + ) + gg, _ = am.descriptor.call_graph( + graph, atype_all, type_embedding=am.descriptor.graph_type_embedding_table() + ) + torch.testing.assert_close(vacuum, gg[3:], rtol=0.0, atol=0.0) + + def test_ragged_and_padded_batches_agree() -> None: """The two layouts are two spellings of one batch, so they must agree. @@ -660,11 +708,12 @@ def test_compact_canonical_eligibility_rejects_other_descriptors() -> None: canonical_model_eligible, ) + learned = SimpleNamespace( + descriptor=SimpleNamespace(compress=True), + fitting_net=object(), + ) model = SimpleNamespace( - atomic_model=SimpleNamespace( - descriptor=SimpleNamespace(compress=True), - fitting_net=object(), - ) + atomic_model=SimpleNamespace(fused_decomposition=lambda: (learned, None)) ) assert not canonical_model_eligible(model) diff --git a/source/tests/pt_expt/model/test_dpa4c_zbl_bridging.py b/source/tests/pt_expt/model/test_dpa4c_zbl_bridging.py new file mode 100644 index 0000000000..3e881fe363 --- /dev/null +++ b/source/tests/pt_expt/model/test_dpa4c_zbl_bridging.py @@ -0,0 +1,575 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""ZBL zone bridging of a DPA4C standard model in pt_expt. + +``bridging_method: zbl`` on a DPA4C model builds the same linear composition +as DPA4 (``LinearEnergyModel`` over ``[learned, InnerPotentialAtomicModel]``). +DPA4C has no message passing and carries the edge envelope in every edge term, +so its bridging window is one inner switch on that envelope, paired with a +clamp of the length the radial features read: a pair at or below the inner +radius leaves the learned model exactly as if its edge were deleted, and the +analytical term supplies the whole pair interaction. +""" + +import copy +import math + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.atomic_model.inner_potential import ( + ELEMENT_TO_Z, +) +from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( + ef_op_available, +) +from deepmd.pt_expt.model.dp_linear_model import ( + LinearEnergyModel, +) +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 deepmd.pt_expt.utils.serialization import ( + _needs_with_comm_artifact, + _resolve_lower_kind, + _trace_and_export, + build_synthetic_graph_inputs, +) + +TYPE_MAP = ["O", "H", "He"] +R_INNER = 0.5 +R_OUTER = 0.8 +PLAIN_CONFIG = { + "type_map": TYPE_MAP, + "descriptor": { + "type": "dpa4c", + "rcut": 4.0, + "channels": 16, + "lmax": 2, + "n_radial": 8, + "radial_modes": 2, + "precision": "float64", + "seed": 7, + }, + "fitting_net": {"neuron": [16, 16], "precision": "float64", "seed": 7}, +} +ZBL_CONFIG = { + **PLAIN_CONFIG, + "bridging_method": "zbl", + "bridging_r_inner": R_INNER, + "bridging_r_outer": R_OUTER, +} +# One O-He pair whose separation is set per test, among atoms that keep every +# other pair beyond the window. +ATYPE = np.array([0, 2, 1, 1]) +DIRECTION = np.array([0.8, 0.36, -0.48]) +ENVIRONMENT = np.array([[-0.4, 0.9, 0.3], [0.2, -0.5, 1.2]]) +#: Magnetic moments of the cluster, in Bohr magnetons. +SPIN = np.array([[0.0, 0.0, 1.4], [0.9, -0.3, 0.2], [-0.5, 1.1, 0.0], [0.3, 0.2, -0.8]]) + + +def _cluster(distance: float) -> np.ndarray: + return np.concatenate([np.zeros((1, 3)), distance * DIRECTION[None], ENVIRONMENT]) + + +def _evaluate(model: torch.nn.Module, coord: np.ndarray) -> dict[str, torch.Tensor]: + """Evaluate one open-boundary frame on the backend device.""" + return model( + torch.tensor( + coord[None], dtype=torch.float64, device=env.DEVICE + ).requires_grad_(True), + torch.tensor(ATYPE[None], dtype=torch.int64, device=env.DEVICE), + box=None, + ) + + +def _analytic_zbl_total(coord: np.ndarray, atype: np.ndarray = ATYPE) -> float: + """Independent reference: direct double loop over open-boundary pairs.""" + a_coeff = (0.18175, 0.50986, 0.28022, 0.028171) + b_coeff = (3.1998, 0.94229, 0.4029, 0.20162) + charge = [float(ELEMENT_TO_Z[TYPE_MAP[t]]) for t in atype] + total = 0.0 + for i in range(len(charge)): + for j in range(i + 1, len(charge)): + r = float(np.linalg.norm(coord[i] - coord[j])) + a = 0.88534 * 0.5291772109 / (charge[i] ** 0.23 + charge[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 * charge[i] * charge[j] / r * phi + return total + + +def _perturb(model: torch.nn.Module) -> None: + """Move the weights off their initialization so outputs depend on geometry.""" + generator = torch.Generator(device="cpu").manual_seed(11) + with torch.no_grad(): + for parameter in model.parameters(): + if parameter.is_floating_point(): + noise = torch.randn( + parameter.shape, + generator=generator, + dtype=parameter.dtype, + device="cpu", + ) + parameter.add_(0.3 * noise.to(parameter.device)) + + +class TestDPA4CZBLBridging: + def setup_method(self) -> None: + self.model = get_model(copy.deepcopy(ZBL_CONFIG)).eval() + _perturb(self.model) + self.learned = self.model.atomic_model.models[0] + + def _plain_twin(self, exclude: list[tuple[int, int]]) -> torch.nn.Module: + """The learned model alone, with the same weights and no window.""" + plain = get_model(copy.deepcopy(PLAIN_CONFIG)).eval() + plain.atomic_model.load_state_dict(self.learned.state_dict()) + plain.atomic_model.descriptor.reinit_exclude(exclude) + return plain + + def test_builder_composes_a_graph_route_linear_model(self) -> None: + assert type(self.model) is LinearEnergyModel + learned, potential = self.model.atomic_model.fused_decomposition() + assert learned is self.learned + assert potential is self.model.atomic_model.models[1].potential + # An explicit window in Å reaches the descriptor as the same pair of + # fractions measured against a unit length scale. + assert learned.descriptor.bridging_scale == "absolute" + assert learned.descriptor.bridging_f_inner == R_INNER + assert learned.descriptor.bridging_f_outer == R_OUTER + assert model_uses_graph_lower(self.model) + # DPA4C exchanges no node features, so bridging needs no with-comm + # artifact (a bridged DPA4 does, for its source gate). + assert _needs_with_comm_artifact(self.model, lower_kind="graph") is False + + @pytest.mark.parametrize("distance", [0.2, 0.45]) + def test_frozen_pair_interacts_through_zbl_alone(self, distance: float) -> None: + coord = _cluster(distance) + energy = _evaluate(self.model, coord)["energy"].item() + removed = _evaluate(self._plain_twin([(0, 2), (2, 0)]), coord)["energy"].item() + visible = _evaluate(self._plain_twin([]), coord)["energy"].item() + # The plain model must see the pair, or the comparison is vacuous. + assert abs(visible - removed) > 1e-3 + assert energy == pytest.approx(removed + _analytic_zbl_total(coord), abs=1e-10) + + def test_pairs_beyond_the_window_add_zbl_to_the_plain_model(self) -> None: + coord = _cluster(R_OUTER + 0.05) + energy = _evaluate(self.model, coord)["energy"].item() + plain = _evaluate(self._plain_twin([]), coord)["energy"].item() + assert energy == pytest.approx(plain + _analytic_zbl_total(coord), abs=1e-10) + + @pytest.mark.parametrize("distance", [0.3, 0.5, 0.65, 0.8, 1.1]) + def test_force_matches_finite_difference(self, distance: float) -> None: + coord = _cluster(distance) + force = _evaluate(self.model, coord)["force"] + step = 1.0e-5 + for axis in range(3): + shift = np.zeros_like(coord) + shift[1, axis] = step + upper = _evaluate(self.model, coord + shift)["energy"].item() + lower = _evaluate(self.model, coord - shift)["energy"].item() + expected = -(upper - lower) / (2.0 * step) + assert force[0, 1, axis].item() == pytest.approx( + expected, rel=1e-6, abs=1e-6 + ) + + def test_serialize_roundtrip(self) -> None: + coord = _cluster(0.65) + data = self.model.serialize() + assert data["type"] == "linear" + restored = LinearEnergyModel.deserialize(data).eval() + torch.testing.assert_close( + _evaluate(restored, coord)["energy"], + _evaluate(self.model, coord)["energy"], + rtol=0.0, + atol=0.0, + ) + + def test_isolated_atoms_keep_the_vacuum_reference(self) -> None: + """A preset bias keeps the learned child referencing the isolated atoms. + + The composition owns the output bias, and its preset reaches the + learned child, whose fitting then still subtracts the network output + of an atom without neighbors. + """ + preset = {"energy": [-430.1, -13.6, -79.0]} + isolated = np.zeros((1, 3)) + energies = [] + for base in (PLAIN_CONFIG, ZBL_CONFIG): + config = copy.deepcopy(base) + config["fitting_net"]["vacuum_ref"] = True + config["preset_out_bias"] = preset + model = get_model(config).eval() + _perturb(model) + fitting = model.atomic_model.fused_decomposition()[0].fitting_net + assert fitting.vacuum_ref + energies.append( + model( + torch.tensor( + isolated[None], dtype=torch.float64, device=env.DEVICE + ), + torch.tensor([[0]], dtype=torch.int64, device=env.DEVICE), + box=None, + )["energy"].item() + ) + # Neither model has computed its statistics, so the reference leaves + # an isolated atom at the zero output bias, up to float64 roundoff + # from the separately evaluated network and vacuum reference. + np.testing.assert_allclose( + energies, [0.0, 0.0], rtol=0.0, atol=8 * np.finfo(np.float64).eps + ) + + def test_declared_tensors_keep_the_adamw_route(self) -> None: + """HybridMuon routes the same descriptor tensors as in the plain model. + + The composition carries the learned model one level deeper, so its + patterns name that path; what must not change is the set of tensors + they reach. + """ + + def routed(model: torch.nn.Module, path: str) -> list[str]: + patterns = model.adam_route_patterns() + assert patterns + names = [ + name + for name, _ in model.named_parameters() + if any(pattern in name.lower() for pattern in patterns) + ] + assert names and all( + name.startswith(f"{path}descriptor.") for name in names + ) + return [name[len(path) :] for name in names] + + plain = get_model(copy.deepcopy(PLAIN_CONFIG)) + assert routed(self.model, "atomic_model.models.0.") == routed( + plain, "atomic_model." + ) + + def test_pair_table_is_rebuilt_not_stored(self) -> None: + potential = self.model.atomic_model.models[1].potential + assert not any("pair_table" in key for key in self.model.state_dict()) + assert potential.pair_table.dtype == torch.float32 + table = potential.pair_table.double().cpu().reshape(4, 4, 8) + # Rows and columns of the padding type vanish. + assert table[3].abs().max() == 0.0 and table[:, 3].abs().max() == 0.0 + radius = 0.7 + series = (table[0, 2, :4] * torch.exp(-table[0, 2, 4:] * radius)).sum() / radius + pair = np.array([[0.0, 0.0, 0.0], [radius, 0.0, 0.0]]) + assert series.item() == pytest.approx( + _analytic_zbl_total(pair, np.array([0, 2])), rel=1e-6 + ) + + +def _native_spin(base: dict) -> dict: + """The same configuration with a magnetic moment on every type.""" + return { + **copy.deepcopy(base), + "spin": {"scheme": "native", "use_spin": [True] * len(TYPE_MAP)}, + } + + +def _evaluate_spin( + model: torch.nn.Module, coord: np.ndarray +) -> dict[str, torch.Tensor]: + """Evaluate one open-boundary frame of a native-spin model.""" + return model( + torch.tensor( + coord[None], dtype=torch.float64, device=env.DEVICE + ).requires_grad_(True), + torch.tensor(ATYPE[None], dtype=torch.int64, device=env.DEVICE), + torch.tensor(SPIN[None], dtype=torch.float64, device=env.DEVICE), + box=None, + ) + + +class TestDPA4CZBLBridgingWithNativeSpin: + """Native spin and ZBL bridging compose on a DPA4C model. + + The analytical term is a function of the separation alone: it adds its + energy and its force to the spin-conditioned model and leaves the magnetic + force untouched. + """ + + def setup_method(self) -> None: + self.model = get_model(_native_spin(ZBL_CONFIG)).eval() + _perturb(self.model) + self.plain = get_model(_native_spin(PLAIN_CONFIG)).eval() + self.plain.atomic_model.load_state_dict( + self.model.atomic_model.models[0].state_dict() + ) + + def test_the_composition_answers_the_spin_capability(self) -> None: + learned, potential = self.model.atomic_model.fused_decomposition() + assert potential is self.model.atomic_model.models[1].potential + assert self.model.atomic_model.supports_native_spin() + # The capability comes from the learned child; the analytical term + # accepts the moments and ignores them. + assert learned.supports_native_spin() + assert not self.model.atomic_model.models[1].supports_native_spin() + + def test_the_analytical_term_leaves_the_magnetic_force_alone(self) -> None: + coord = _cluster(R_OUTER + 0.05) + bridged = _evaluate_spin(self.model, coord) + plain = _evaluate_spin(self.plain, coord) + assert bridged["energy"].item() == pytest.approx( + plain["energy"].item() + _analytic_zbl_total(coord), abs=1e-10 + ) + torch.testing.assert_close( + bridged["force_mag"], plain["force_mag"], rtol=0.0, atol=0.0 + ) + # The analytical force must reach the atoms, or the comparison above + # holds for a term that does nothing. + assert (bridged["force"] - plain["force"]).abs().max().item() > 1e-3 + + +def test_children_reference_the_isolated_atoms_by_the_composition_preset() -> None: + """A child's fitting references the isolated atoms by the composition preset. + + The composition computes the output bias, so a preset in the configuration + of a child alone fixes nothing and leaves the reference off. + """ + preset = {"energy": [-430.1, -13.6, -79.0]} + child = { + "descriptor": copy.deepcopy(PLAIN_CONFIG["descriptor"]), + "fitting_net": { + **copy.deepcopy(PLAIN_CONFIG["fitting_net"]), + "vacuum_ref": True, + }, + "preset_out_bias": preset, + } + config = { + "type": "linear_ener", + "type_map": TYPE_MAP, + "models": [child, copy.deepcopy(child)], + "weights": "mean", + } + model = get_model(copy.deepcopy(config)) + assert not any(sub.fitting_net.vacuum_ref for sub in model.atomic_model.models) + config["preset_out_bias"] = preset + model = get_model(config) + assert all(sub.fitting_net.vacuum_ref for sub in model.atomic_model.models) + + +def _compressed_config(window: tuple[float, float] | None) -> dict: + """A float32 DPA4C model the fused operators can serve.""" + config = { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa4c", + "rcut": 3.0, + "channels": 16, + "lmax": 3, + "n_radial": 8, + "radial_modes": 2, + "precision": "float32", + "seed": 17, + }, + "fitting_net": { + "type": "ener", + "neuron": [32, 32], + "activation_function": "silu", + "resnet_dt": False, + "precision": "float32", + "seed": 19, + }, + } + if window is not None: + config.update( + bridging_method="zbl", + bridging_r_inner=window[0], + bridging_r_outer=window[1], + ) + return config + + +class TestCompressedDPA4CZBLBridging: + """A bridged compressed model keeps the fused and the canonical routes.""" + + def setup_method(self) -> None: + # The window follows the sample, so that its edges populate the frozen + # zone, the transition zone and the plain zone alike. + plain = get_model(_compressed_config(None)).to(env.DEVICE).eval() + self.sample = build_synthetic_graph_inputs( + plain, + e_max=None, + nframes=2, + nloc=9, + dtype=torch.float32, + device=env.DEVICE, + ) + edge_vec, edge_mask = self.sample[4], self.sample[5] + length = edge_vec[edge_mask].norm(dim=-1) + window = tuple( + float(value) + for value in torch.quantile( + length, + torch.tensor([0.2, 0.5], dtype=length.dtype, device=length.device), + ) + ) + self.model = get_model(_compressed_config(window)).to(env.DEVICE).eval() + _perturb(self.model) + self.model.atomic_model.enable_compression(0.5, 1.0, 0.002, 0.002) + assert ((length > window[0]) & (length < window[1])).sum() > 4 + assert (length < window[0]).any() + + def _lower(self, n_local: torch.Tensor | None = None) -> dict[str, torch.Tensor]: + ( + atype, + n_node, + sample_n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + ) = self.sample + return self.model.forward_common_lower_graph( + atype, + n_node, + sample_n_local if n_local is None else n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + destination_sorted=True, + do_atomic_virial=True, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + ) + + def _autograd_lower( + self, + monkeypatch: pytest.MonkeyPatch, + n_local: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """The portable composition: learned descriptor plus the ZBL sibling.""" + from deepmd.pt_expt.model import ( + make_model, + ) + + with monkeypatch.context() as patch: + patch.setattr(make_model, "fused_energy_force_enabled", lambda: False) + return self._lower(n_local) + + @pytest.mark.skipif( + not ef_op_available(), + reason="the fused DPA4C energy-force operators are required", + ) + @pytest.mark.parametrize("ghosts", [0, 3]) + def test_fused_route_matches_the_autograd_composition( + self, + monkeypatch: pytest.MonkeyPatch, + ghosts: int, + ) -> None: + monkeypatch.setenv("DP_CUDA_INFER", "2") + n_local = self.sample[1] - ghosts + reference = self._autograd_lower(monkeypatch, n_local) + fused = self._lower(n_local) + # The analytical term must matter, or the comparison is vacuous. + learned = self.model.atomic_model.models[0] + zbl = reference["energy_redu"].detach().sum() - ( + learned.forward_common_atomic_graph(*self._graph_and_types())["energy"] + .detach() + .sum() + ) + assert ghosts or abs(float(zbl)) > 1.0 + # A silent fallback to the autograd lower would compare it with itself. + assert not torch.equal(fused["energy_derv_r"], reference["energy_derv_r"]) + for key in ( + "energy", + "energy_redu", + "energy_derv_r", + "energy_derv_c", + "energy_derv_c_redu", + ): + torch.testing.assert_close( + fused[key].double(), + reference[key].double(), + atol=2e-4, + rtol=2e-4, + ) + + def _graph_and_types(self) -> tuple: + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + ) + + atype, n_node, n_local, edge_index, edge_vec, edge_mask = self.sample[:6] + graph = NeighborGraph( + n_node=n_node, + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, + n_local=n_local, + ) + return graph, atype + + @pytest.mark.skipif( + env.DEVICE.type != "cuda", reason="the compact canonical route is CUDA only" + ) + @pytest.mark.parametrize("tile", [None, "4"]) + def test_canonical_route_matches_the_autograd_composition( + self, + monkeypatch: pytest.MonkeyPatch, + tile: str | None, + ) -> None: + if tile is not None: + monkeypatch.setenv("DP_NODE_TILE", tile) + reference = self._autograd_lower(monkeypatch) + atype, n_node, n_local, edge_index, edge_vec = self.sample[:5] + destination_row_ptr, source_order, source_row_ptr = self.sample[7:10] + physical = int(destination_row_ptr[-1]) + canonical = self.model.forward_lower_canonical_graph( + atype, + n_node, + n_local, + edge_index[0][:physical].to(torch.uint32).contiguous(), + edge_vec[:physical].contiguous(), + destination_row_ptr, + source_row_ptr, + source_order[:physical].to(torch.uint32).contiguous(), + do_atomic_virial=True, + ) + for key, reference_key in ( + ("atom_energy", "energy"), + ("energy", "energy_redu"), + ("force", "energy_derv_r"), + ("atom_virial", "energy_derv_c"), + ("virial", "energy_derv_c_redu"), + ): + torch.testing.assert_close( + canonical[key].double().reshape(-1), + reference[reference_key].double().reshape(-1), + atol=2e-4, + rtol=2e-4, + ) + + def test_compressed_model_exports_like_its_learned_part(self) -> None: + data = {"model": self.model.to("cpu").serialize()} + expected = "dpa4c_canonical" if env.DEVICE.type == "cuda" else "graph" + assert _resolve_lower_kind("model.pt2", data, "auto") == expected + exported, metadata, _model_json, _output_keys = _trace_and_export( + data, + lower_kind="graph", + do_atomic_virial=True, + ) + assert isinstance(exported, torch.export.ExportedProgram) + assert metadata["graph_edge_dtype"] == "float32" + assert metadata["has_comm_artifact"] is False diff --git a/source/tests/pt_expt/model/test_get_model_bridging.py b/source/tests/pt_expt/model/test_get_model_bridging.py index e4b9874519..ecead9a672 100644 --- a/source/tests/pt_expt/model/test_get_model_bridging.py +++ b/source/tests/pt_expt/model/test_get_model_bridging.py @@ -179,6 +179,26 @@ def test_compile_attention_probe_tolerates_composition() -> None: _warn_compiled_attention(model, "Default") +def test_fine_tune_locates_the_learned_part_by_structure() -> None: + """Fine-tuning finds the learned child of a composition by its descriptor. + + A descriptor exclusion gives the child its own pair exclusion, which keeps + the composition off the fused route, but not its learned part off its path. + """ + from deepmd.pt_expt.train.training import ( + _learned_part_path, + ) + + plain = _dpa4_standard_config() + plain["type"] = "dpa4" + plain["descriptor"]["exclude_types"] = [[0, 1]] + bridged = _bridged(copy.deepcopy(plain)) + bridged_model = get_model(bridged) + assert bridged_model.atomic_model.fused_decomposition() is None + assert _learned_part_path(get_model(plain)) == "atomic_model." + assert _learned_part_path(bridged_model) == "atomic_model.models.0." + + def _canonical_config() -> dict: """The bridged config spelled canonically (issue #5948).""" base = _dpa4_standard_config() @@ -224,7 +244,9 @@ def test_canonical_composition_builds() -> None: assert len(model.atomic_model.models) == 2 learned = model.atomic_model.models[0] assert learned.descriptor.bridging_switch is not None - assert float(learned.descriptor.inner_clamp.r_inner) == 0.8 + assert float(learned.descriptor.bridging_switch.f_inner) == 0.8 + assert float(learned.descriptor.bridging_clamp.f_freeze) == pytest.approx(0.96) + assert learned.descriptor.bridging_scale == "absolute" def test_canonical_matches_sugar_serialize() -> None: diff --git a/source/tests/pt_expt/model/test_linear_model.py b/source/tests/pt_expt/model/test_linear_model.py index 8c95f3e69e..ecc2bd9a31 100644 --- a/source/tests/pt_expt/model/test_linear_model.py +++ b/source/tests/pt_expt/model/test_linear_model.py @@ -141,6 +141,24 @@ def _prepare_lower_inputs(self): torch.tensor(mapping, dtype=torch.int64, device=self.device), ) + def test_integer_mappings_are_buffers(self) -> None: + """The type mappings of the children are buffers, not parameters. + + Every parameter is then floating point, which a sharded optimizer such + as ``ZeroRedundancyOptimizer`` requires, while the state-dict keys of + the mappings stay those of a list attribute. + """ + md_pt = LinearEnergyModel.deserialize( + self._make_dp_linear_model().serialize() + ).to(self.device) + self.assertTrue(all(p.is_floating_point() for p in md_pt.parameters())) + mapping_list = md_pt.atomic_model.mapping_list + self.assertEqual(len(mapping_list), 2) + self.assertTrue(all(m.dtype == torch.int64 for m in mapping_list)) + state = md_pt.state_dict() + self.assertIn("atomic_model.mapping_list.0", state) + self.assertIn("atomic_model.mapping_list.1", state) + def test_linear_model_consistency(self) -> None: """Create a LinearEnergyModel, run forward() and forward_lower(), verify outputs have correct keys and shapes. diff --git a/source/tests/pt_expt/model/test_zbl_bridging.py b/source/tests/pt_expt/model/test_zbl_bridging.py index 28ec95b819..b26bdc9234 100644 --- a/source/tests/pt_expt/model/test_zbl_bridging.py +++ b/source/tests/pt_expt/model/test_zbl_bridging.py @@ -186,6 +186,297 @@ def test_zbl_child_adds_positive_energy(self) -> None: ) assert diff > 1e-3 + def _learned_energies_on_graph(self, backend: str, graph, atype) -> torch.Tensor: + """Atomic energies of the learned model alone on an explicit edge set.""" + if backend == "pt": + model = copy.deepcopy(self.pt_model) + model.inter_potential = None + coord = torch.zeros((1, atype.shape[1], 3), dtype=torch.float64) + out = model.forward_lower( + coord, + atype, + graph.edge_index, + graph.edge_vec, + graph.edge_index, + graph.edge_mask, + ) + return out["atom_energy"].detach().reshape(-1) + learned = self.pt_expt_model.atomic_model.models[0] + out = learned.forward_atomic_graph(graph, atype.reshape(-1)) + return out["energy"].detach().reshape(-1) + + def _learned_descriptor_on_graph(self, backend: str, graph, atype) -> torch.Tensor: + """Descriptor of the learned model alone on an explicit edge set.""" + descriptor = self._descriptor(backend) + if backend == "pt": + out = descriptor.forward_with_edges( + extended_coord=torch.zeros((1, atype.shape[1], 3), dtype=torch.float64), + extended_atype=atype, + edge_index=graph.edge_index, + edge_vec=graph.edge_vec, + edge_mask=graph.edge_mask, + )[0] + else: + out = descriptor.call_graph(graph, atype.reshape(-1))[0] + return out.detach().reshape(atype.shape[1], -1) + + def _descriptor(self, backend: str): + """The descriptor of the learned model of either backend.""" + if backend == "pt": + return self.pt_model.atomic_model.descriptor + return self.pt_expt_model.atomic_model.models[0].descriptor + + def _bias(self, backend: str, atype: torch.Tensor) -> torch.Tensor: + """The per-atom end point of the fade: the fitting bias of each atom's type.""" + fitting = ( + self.pt_model.atomic_model.fitting_net + if backend == "pt" + else self.pt_expt_model.atomic_model.models[0].fitting_net + ) + bias = torch.as_tensor(fitting.bias_atom_e, dtype=torch.float64).reshape(-1) + return bias[atype.reshape(-1)] + + @staticmethod + def _quartet(gap: float = 0.6, shift=(0.0, 0.0, 0.0)): + """Four atoms ``(A, B, C, D)`` with ``A`` and ``B`` a frozen pair. + + ``A`` and ``B`` sit ``gap`` apart, inside the frozen zone of the 0.8 / + 1.2 Å window, and ``shift`` translates that pair rigidly. ``C`` and + ``D`` are ordinary atoms: every one of their separations from the pair + and from each other is beyond the outer radius and within the cutoff, + so they carry an ordinary bond to each other. + """ + sx, sy, sz = shift + coord = torch.tensor( + [ + [ + [sx, sy, sz], + [sx + gap, sy, sz], + [2.4, 0.3, 0.1], + [-1.0, 1.2, -0.4], + ] + ], + dtype=torch.float64, + ) + atype = torch.tensor([[0, 1, 0, 1]], dtype=torch.int64) + return coord, atype + + @classmethod + def _quartet_graph(cls, gap: float = 0.6, shift=(0.0, 0.0, 0.0)): + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, + ) + + coord, atype = cls._quartet(gap=gap, shift=shift) + graph = build_neighbor_graph( + coord, atype, None, ZBL_CONFIG["descriptor"]["rcut"], with_csr=True + ) + return graph, atype + + @staticmethod + def _mask_edges(graph, sources, destinations): + """The same graph with every ``sources -> destinations`` edge masked.""" + import dataclasses + + src, dst = graph.edge_index[0], graph.edge_index[1] + drop = torch.isin(src, torch.tensor(sources)) & torch.isin( + dst, torch.tensor(destinations) + ) + return ( + dataclasses.replace(graph, edge_mask=graph.edge_mask & ~drop), + int((drop & graph.edge_mask).sum()), + ) + + @pytest.mark.parametrize("backend", ["pt", "pt_expt"]) + def test_a_frozen_pair_leaves_the_environment_of_the_others(self, backend) -> None: + """The four edges out of the frozen pair into ``C`` and ``D`` are gone. + + ``A`` and ``B`` are closer than the inner radius, so the leave-one-out + gate is zero on every edge they emit except the two between them. + Deleting those four edges by hand must therefore change no learned + energy at all: the gate is folded into the envelope, the one factor + the messages, the environment seed, the attention masses and the + degree normalization all weight an edge by. The frozen atoms + themselves are left with the per-type bias the fitting stores, the end + point of the fade. + """ + graph, atype = self._quartet_graph() + pruned, dropped = self._mask_edges(graph, sources=[0, 1], destinations=[2, 3]) + # Twelve real edges: six pairs, of which four edges leave the pair. + assert int(graph.edge_mask.sum()) == 12 + assert dropped == 4 + + full = self._learned_energies_on_graph(backend, graph, atype) + deleted = self._learned_energies_on_graph(backend, pruned, atype) + torch.testing.assert_close(full, deleted, rtol=0.0, atol=1e-12) + + bias = self._bias(backend, atype) + torch.testing.assert_close(full[:2], bias[:2], rtol=0.0, atol=1e-12) + # the ordinary atoms still carry a learned energy of their own + assert (full[2:] - bias[2:]).abs().min().item() > 1e-6 + + @pytest.mark.parametrize("backend", ["pt", "pt_expt"]) + def test_the_edge_inside_a_frozen_pair_stays_alive(self, backend) -> None: + """``A`` and ``B`` go on seeing each other at the clamped distance. + + The gate an edge carries is the product over the other pairs of its + source, so the pair ``(A, B)`` leaves its own two edges untouched. + Masking them is therefore a visible change to the descriptors of ``A`` + and ``B``, which is what separates the leave-one-out gate from a gate + that folded a node's whole edge set. It leaves the energies of the + ordinary atoms alone, since the pair had already vanished from their + environment. + + The mask is also what carries the freeze: a masked edge contributes the + multiplicative identity to the products, so removing the pair edge + removes the amplitude of zero with it and the two atoms stop being + faded. Their energies sit on the per-type bias while the pair edge is + there and leave it once the edge is gone, which pins the mask + convention alongside. + """ + graph, atype = self._quartet_graph() + pruned, _ = self._mask_edges(graph, sources=[0, 1], destinations=[2, 3]) + isolated, dropped = self._mask_edges( + pruned, sources=[0, 1], destinations=[0, 1] + ) + assert dropped == 2 + + with_pair = self._learned_descriptor_on_graph(backend, pruned, atype) + without_pair = self._learned_descriptor_on_graph(backend, isolated, atype) + assert (with_pair[:2] - without_pair[:2]).abs().max().item() > 1e-6 + + energy_with_pair = self._learned_energies_on_graph(backend, pruned, atype) + energy_without_pair = self._learned_energies_on_graph(backend, isolated, atype) + torch.testing.assert_close( + energy_with_pair[2:], energy_without_pair[2:], rtol=0.0, atol=1e-12 + ) + bias = self._bias(backend, atype) + torch.testing.assert_close(energy_with_pair[:2], bias[:2], rtol=0.0, atol=1e-12) + assert (energy_without_pair[:2] - bias[:2]).abs().max().item() > 1e-6 + + @pytest.mark.parametrize("backend", ["pt", "pt_expt"]) + @pytest.mark.parametrize( + "gap, shift", + [ + (0.6, (0.1, 0.05, 0.0)), # rigid translation of the frozen pair + (0.7, (0.0, 0.0, 0.0)), # a different separation, still frozen + ], + ) + def test_a_frozen_pair_moves_without_changing_any_energy( + self, backend, gap, shift + ) -> None: + """Nothing the frozen pair does inside its own zone reaches an energy. + + Both motions change the geometry every edge of the pair carries, and + both leave every learned energy exactly where it was: the atoms of the + pair are pinned to their per-type bias and the ordinary atoms no + longer see the pair at all. + """ + reference = self._learned_energies_on_graph(backend, *self._quartet_graph()) + moved = self._learned_energies_on_graph( + backend, *self._quartet_graph(gap=gap, shift=shift) + ) + torch.testing.assert_close(reference, moved, rtol=0.0, atol=1e-12) + + @pytest.mark.parametrize("backend", ["pt", "pt_expt"]) + def test_clearing_the_gate_reopens_the_leak(self, backend) -> None: + """Ablation: without the switch the frozen pair is back in the graph. + + Clearing ``bridging_switch`` leaves the distance clamp in place, so a + difference here pins the gate, and not the clamp, as the owner of the + isolation above. + """ + graph, atype = self._quartet_graph() + pruned, _ = self._mask_edges(graph, sources=[0, 1], destinations=[2, 3]) + deleted = self._learned_energies_on_graph(backend, pruned, atype) + + descriptor = self._descriptor(backend) + switch = descriptor.bridging_switch + descriptor.bridging_switch = None + try: + ungated = self._learned_energies_on_graph(backend, graph, atype) + finally: + descriptor.bridging_switch = switch + assert (ungated - deleted).abs().max().item() > 1e-6 + + @pytest.mark.parametrize("backend", ["pt", "pt_expt"]) + def test_a_closing_dimer_follows_the_switch(self, backend) -> None: + """A dimer's learned energy is proportional to the switch amplitude. + + On a two-atom system the leave-one-out gate is the empty product, so + both edges keep their full envelope and the descriptor reads exactly + the distance the clamp displays; the per-node gate is the switching + amplitude of the single pair. The clamp freezes that displayed distance + two fifths into the window for every separation at or below that + point, so from the inner radius to the freeze point the descriptor + never changes and + + ``E(r) - bias = w(r) * (E_unbridged(r_z) - bias)`` + + holds with one fixed factor: the deviation from the bias is the + amplitude times a constant, to round-off. Below the inner radius the + amplitude is zero and only the bias is left. + """ + from deepmd.dpmodel.descriptor.dpa4_nn.radial import ( + BridgingSwitch, + InnerClamp, + ) + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, + ) + + f_inner = ZBL_CONFIG["bridging_r_inner"] + f_outer = ZBL_CONFIG["bridging_r_outer"] + # The absolute scale gives every pair the unit length scale, so the two + # radii are read straight off the separation in Å. + contact = np.ones((1,)) + switch = BridgingSwitch(f_inner, f_outer) + clamp = InnerClamp(f_inner, f_outer) + freeze = f_inner + 0.4 * (f_outer - f_inner) + assert clamp.f_freeze == pytest.approx(freeze) + atype = torch.tensor([[0, 1]], dtype=torch.int64) + bias = self._bias(backend, atype) + descriptor = self._descriptor(backend) + + def _dimer_energies(separation: float, bridged: bool) -> torch.Tensor: + coord = torch.tensor( + [[[0.0, 0.0, 0.0], [separation, 0.0, 0.0]]], dtype=torch.float64 + ) + graph = build_neighbor_graph( + coord, atype, None, ZBL_CONFIG["descriptor"]["rcut"], with_csr=True + ) + if bridged: + return self._learned_energies_on_graph(backend, graph, atype) + saved = descriptor.bridging_clamp, descriptor.bridging_switch + descriptor.bridging_clamp = None + descriptor.bridging_switch = None + try: + return self._learned_energies_on_graph(backend, graph, atype) + finally: + descriptor.bridging_clamp, descriptor.bridging_switch = saved + + # below the inner radius the pair contributes its bias and nothing else + for separation in (0.40, 0.60, 0.79, f_inner): + torch.testing.assert_close( + _dimer_energies(separation, bridged=True), bias, rtol=0.0, atol=1e-12 + ) + + # the one fixed factor: the unbridged deviation at the frozen distance + factor = _dimer_energies(freeze, bridged=False) - bias + # the fixture must not make the proportionality vacuous + assert factor.abs().max().item() > 1e-4 + for separation in (0.85, 0.90, 0.95, freeze): + displayed = float(clamp.call(np.array([separation]), contact)[0]) + assert displayed == pytest.approx(freeze, abs=1e-15) + amplitude = float(switch.call(np.array([separation]), contact)[0]) + assert 0.0 < amplitude < 1.0 + torch.testing.assert_close( + _dimer_energies(separation, bridged=True) - bias, + amplitude * factor, + rtol=0.0, + atol=1e-12, + ) + def test_force_matches_finite_difference(self) -> None: """F = -dE/dx through the shared-edge-leaf summed autograd.""" eps = 1e-5 @@ -203,6 +494,36 @@ def test_force_matches_finite_difference(self) -> None: float(force[atom, comp]), fd, rtol=1e-6, atol=1e-6 ) + @pytest.mark.parametrize( + "version", + [None, 1.0, 1.1, 1.2], # stored descriptor version; None = untagged checkpoint + ) + def test_legacy_bridged_state_is_rejected(self, version: float | None) -> None: + """A composition cannot restore weights trained under an older window.""" + state = self.pt_expt_model.state_dict() + version_key = "atomic_model.models.0.descriptor.version_tensor" + if version is None: + del state[version_key] + else: + state[version_key] = torch.full_like(state[version_key], version) + with pytest.raises(ValueError, match="Retrain the bridged model"): + self.pt_expt_model.load_state_dict(state) + + def test_legacy_plain_state_into_bridged_model_requires_finetune(self) -> None: + """The learned child applies the same target-based state-loading guard.""" + config = copy.deepcopy(ZBL_CONFIG) + config["bridging_method"] = "none" + plain = get_model(config) + state = plain.atomic_model.state_dict() + version_key = "descriptor.version_tensor" + state[version_key] = torch.full_like(state[version_key], 1.2) + learned = self.pt_expt_model.atomic_model.models[0] + with pytest.raises( + ValueError, + match=r"format version 1\.2 into a bridged descriptor.*--finetune", + ): + learned.load_state_dict(state) + def test_serialize_roundtrip(self) -> None: from deepmd.pt_expt.model.model import ( BaseModel, @@ -293,8 +614,11 @@ def test_construction_composes_and_keeps_spin(self) -> None: assert model.has_spin() is True kinds = [type(c).__name__ for c in model.atomic_model.models] assert kinds[1] == "InnerPotentialAtomicModel", kinds - # bridging radii still reach the LEARNED child's descriptor - assert float(model.atomic_model.models[0].descriptor.inner_clamp.r_inner) == 0.8 + # the bridging window still reaches the LEARNED child's descriptor, + # an explicit radius in Å arriving as a fraction of a unit length scale + learned = model.atomic_model.models[0] + assert float(learned.descriptor.bridging_switch.f_inner) == 0.8 + assert float(learned.descriptor.bridging_clamp.f_freeze) == pytest.approx(0.96) def test_forward_energy_force_force_mag(self) -> None: model = get_model(_native_spin_zbl_config()).to(torch.device("cpu")).eval() @@ -378,16 +702,21 @@ def test_native_spin_with_bridging_dpmodel() -> None: ] -def test_bridging_radii_defaults() -> None: - """bridging_r_inner/r_outer default to 0.5/0.8 on the learned child.""" +def test_bridging_window_defaults() -> None: + """Without explicit radii the window defaults to fractions of the bond length. + + The learned child then measures the window against each pair's own length + scale, so the switch carries the two dimensionless bounds and the clamp + its freeze point two fifths between them, rather than a distance in Å. + """ 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 + descriptor = model.atomic_model.models[0].descriptor + assert float(descriptor.bridging_switch.f_inner) == 0.26 + assert float(descriptor.bridging_switch.f_outer) == 0.80 + assert float(descriptor.bridging_clamp.f_freeze) == pytest.approx(0.476) class TestZBLBridgingExportAndTraining: @@ -516,8 +845,8 @@ class TestInnerPotentialChangeTypeMapPtExpt: Exercised through the REAL composition: ``LinearEnergyModel`` -> ``LinearEnergyAtomicModel`` -> ``InnerPotentialAtomicModel`` -> - ``InnerPotential``. Inside a pt_expt module tree the element lookup is a - wrapped torch buffer, so the rebuild must land on the same + ``InnerPotential``. Inside a pt_expt module tree the coefficient tables + are wrapped torch buffers, so the rebuild must land on the same device/namespace (review 3649295675) -- a numpy rebuild would desync the buffer or fail outright on CUDA. """ @@ -558,41 +887,44 @@ def _build(self, 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.""" + """Precondition: inside pt_expt the tables are torch buffers.""" 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; " + table = self._zbl_child(self._build(["Ni", "O"])).potential.series_table + assert isinstance(table, torch.Tensor), ( + "the pt_expt wrapper no longer converts the table to a tensor; " "this test would stop covering the device-safe rebuild" ) - assert z.device.type == torch.device(_env.DEVICE).type + assert table.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. + # limitation. The tables under test belong 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 + # anti-vacuity: Ni-Ni and O-O must be far apart, else a stale table # 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] + torch.testing.assert_close( + child.potential.series_table, fresh.potential.series_table + ) 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] + table = child.potential.series_table + assert isinstance(table, torch.Tensor), "the rebuild dropped out of torch" + assert table.device.type == torch.device(_env.DEVICE).type + # one row per ordered pair of the three types plus the padding type + assert tuple(table.shape) == (16, 8) # the new type is addressable -- a stale (length-2) table raises here np.testing.assert_allclose( self._pair_energy(child, atype_value=2), diff --git a/source/tests/pt_expt/test_finetune.py b/source/tests/pt_expt/test_finetune.py index 7413743481..859e1df8c2 100644 --- a/source/tests/pt_expt/test_finetune.py +++ b/source/tests/pt_expt/test_finetune.py @@ -771,6 +771,103 @@ def test_finetune_random_fitting(self) -> None: os.chdir(old_cwd) shutil.rmtree(tmpdir, ignore_errors=True) + def _finetune_across_bridging( + self, + source_bridged: bool, + model_branch: str = "", + ) -> None: + """Fine-tune between a plain DPA4C model and its ZBL-bridged composition. + + The composition holds the learned descriptor and fitting net as its + first sub-model, so the weights of that part come from the checkpoint + although their path differs; a RANDOM fitting takes the descriptor + alone and computes the output statistics of the new fitting. + """ + plain = { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa4c", + "rcut": 4.0, + "channels": 8, + "lmax": 2, + "n_radial": 8, + "seed": 1, + }, + "fitting_net": {"neuron": [16, 16], "seed": 1}, + "data_stat_nbatch": 1, + } + bridged = { + **deepcopy(plain), + "bridging_method": "zbl", + "bridging_r_inner": 0.5, + "bridging_r_outer": 0.8, + } + source, target = (bridged, plain) if source_bridged else (plain, bridged) + source_path, target_path = ( + (".atomic_model.models.0.", ".atomic_model.") + if source_bridged + else (".atomic_model.", ".atomic_model.models.0.") + ) + tmpdir = tempfile.mkdtemp(prefix="pt_expt_ft_bridged_") + old_cwd = os.getcwd() + os.chdir(tmpdir) + try: + config = normalize( + update_deepmd_input( + _make_config(self.data_dir, source, numb_steps=1), warning=False + ) + ) + ckpt_path = self._train_pretrained(config, tmpdir) + + ft_config = normalize( + update_deepmd_input( + _make_config(self.data_dir, target, numb_steps=1), warning=False + ) + ) + ft_config["model"], finetune_links = get_finetune_rules( + ckpt_path, ft_config["model"], model_branch=model_branch + ) + trainer_ft = get_trainer( + ft_config, + finetune_model=ckpt_path, + finetune_links=finetune_links, + ) + + pretrained_state = torch.load( + ckpt_path, map_location=DEVICE, weights_only=True + )["model"] + ft_state = trainer_ft.wrapper.state_dict() + parts = ["descriptor."] + if model_branch != "RANDOM": + parts.append("fitting_net.") + learned = [ + key + for key in ft_state + if any(f"{target_path}{part}" in key for part in parts) + ] + self.assertTrue(learned) + for key in learned: + torch.testing.assert_close( + ft_state[key], + pretrained_state[key.replace(target_path, source_path)], + msg=f"{key} should come from the checkpoint", + ) + finally: + os.chdir(old_cwd) + shutil.rmtree(tmpdir, ignore_errors=True) + + def test_finetune_plain_into_bridged(self) -> None: + """A plain DPA4C checkpoint fine-tunes into its ZBL-bridged composition.""" + self._finetune_across_bridging(source_bridged=False) + + def test_finetune_bridged_into_plain(self) -> None: + """A ZBL-bridged DPA4C checkpoint fine-tunes into the plain model.""" + self._finetune_across_bridging(source_bridged=True) + + def test_finetune_random_fitting_into_bridged(self) -> None: + """A RANDOM fitting of a bridged target takes the plain descriptor.""" + self._finetune_across_bridging(source_bridged=False, model_branch="RANDOM") + def test_finetune_from_pte(self) -> None: """Train -> freeze to .pte -> finetune from .pte -> verify checkpoint.""" from deepmd.pt_expt.entrypoints.main import ( diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py index fc7a64e7ee..539b72f2c8 100644 --- a/source/tests/pt_expt/test_lmdb_training.py +++ b/source/tests/pt_expt/test_lmdb_training.py @@ -9,6 +9,7 @@ :class:`LmdbDataSystem` and runs a few training steps. """ +import copy import os import shutil import tempfile @@ -26,6 +27,10 @@ normalize_batch, split_batch, ) +from deepmd.dpmodel.utils.dist_check import ( + compute_min_pair_margin_batch, + pair_half_thresholds, +) from deepmd.dpmodel.utils.lmdb_data import ( collate_lmdb_frames, ) @@ -391,6 +396,9 @@ def test_distributed_batches_are_sharded_with_equal_epoch_lengths(self) -> None: try: self.assertEqual([system.nbatches for system in systems], [[3], [3]]) self.assertEqual([len(system._sampler) for system in systems], [2, 2]) + self.assertEqual( + [system.get_batch_pass_length() for system in systems], [2, 2] + ) batches = [ [system.get_batch()["fid"] for _ in range(len(system._sampler))] for system in systems @@ -408,6 +416,62 @@ def test_distributed_batches_are_sharded_with_equal_epoch_lengths(self) -> None: } self.assertEqual(observed, set(range(8))) + def test_ragged_filter_uses_the_consumed_pass_length(self) -> None: + """Prefetching a shorter pass does not hide the current pass's last frame.""" + from ..common.dpmodel.test_lmdb_data import ( + _create_mixed_nloc_lmdb, + ) + + path = _create_mixed_nloc_lmdb(os.path.join(self.tmpdir, "filter-pass.lmdb")) + with lmdb.open(path) as database, database.begin(write=True) as txn: + for index in range(10): + key = format(index, "012d").encode() + frame = msgpack.unpackb(txn.get(key), raw=False) + natoms = frame["coords"]["shape"][0] + coord = np.zeros((natoms, 3)) + if index == 3: + coord[:, 0] = np.arange(natoms) + frame["coords"] = _encode_array(coord) + frame["cells"] = _encode_array(20.0 * np.eye(3)) + txn.put(key, msgpack.packb(frame, use_bin_type=True)) + + for num_workers in (0, 2): + with self.subTest(num_workers=num_workers): + data = LmdbDataSystem( + path, ["O", "H"], "mix:24", seed=1, num_workers=num_workers + ) + data.use_ragged_batches(True) + data.add_data_requirements( + [ + DataRequirementItem( + "pair_margin", 1, default=0.5, source_policy="derived" + ) + ] + ) + trainer = Trainer.__new__(Trainer) + trainer.multi_task = False + trainer.world_size = 1 + trainer.training_data_by_task = {"Default": data} + trainer.min_pair_dist_by_task = {"Default": 0.5} + try: + # This shuffle packs five batches and leaves frame 3 last; + # the following shuffle packs four batches. + self.assertEqual(data.get_batch_pass_length(), 5) + with patch.object( + data, "get_batch", wraps=data.get_batch + ) as get_batch: + _, labels = trainer.get_data(is_train=True) + self.assertEqual(get_batch.call_count, 5) + self.assertEqual(labels["pair_margin"].numel(), 1) + self.assertGreaterEqual(float(labels["pair_margin"].item()), 1.0) + self.assertEqual(data.get_batch_pass_length(), 5) + data.get_batch() + self.assertEqual(data.get_batch_pass_length(), 4) + if num_workers > 1: + self.assertTrue(data._batch_iterator.started) + finally: + data.close() + def test_add_data_requirements_passthrough(self) -> None: ds = LmdbDataSystem( lmdb_path=self.lmdb_path, @@ -1023,6 +1087,50 @@ def _run_spin(self, model: dict): finally: os.chdir(cwd) + def test_min_pair_dist_filters_ragged_batches(self) -> None: + """Too-close frames leave a ragged batch together with their atoms.""" + config = self._config(self._dpa1()) + config["training"]["training_data"]["min_pair_dist"] = 1.0e-3 + cwd = os.getcwd() + os.chdir(self.tmpdir) + try: + # A frame that clears its window carries a margin it reaches + # rather than its minimum, so the distances that set a splitting + # threshold are measured on the geometry itself. A unit pair + # threshold makes each margin the minimum pair distance in Å. + batch = get_trainer(copy.deepcopy(config)).training_data.get_batch() + distances = compute_min_pair_margin_batch( + np.asarray(batch["coord"]), + np.asarray(batch["box"]), + np.asarray(batch["atype"]), + pair_half_thresholds(1.0, "absolute", ntypes=2), + n_node=np.asarray(batch["n_node"]), + ) + threshold = float(np.median(distances)) + self.assertTrue((distances < threshold).any()) + + config["training"]["training_data"]["min_pair_dist"] = threshold + trainer = get_trainer(copy.deepcopy(config)) + self.assertTrue(trainer.training_data._reader.ragged_batches) + dropped = False + for _ in range(6): + input_dict, label_dict = trainer.get_data(is_train=True) + kept = label_dict["pair_margin"].reshape(-1) + nframes = input_dict["n_node"].shape[0] + nnodes = int(input_dict["n_node"].sum()) + self.assertTrue(bool((kept >= 1.0).all())) + self.assertEqual(kept.numel(), nframes) + self.assertEqual(label_dict["energy"].shape[0], nframes) + self.assertEqual(input_dict["coord"].shape[0], nnodes) + self.assertEqual(input_dict["atype"].shape[0], nnodes) + self.assertEqual(label_dict["force"].shape[0], nnodes) + # `mix:27` packs three frames unless the filter removes one. + dropped = dropped or nframes < 3 + self.assertTrue(dropped) + trainer.run() + finally: + os.chdir(cwd) + def test_graph_model_trains_on_a_flat_node_axis(self) -> None: """DPA1 reads a graph lower, so its batches carry no padded row.""" trainer, batch = self._run(self._dpa1()) diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py index 4fdad73dc8..fa4199195b 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -30,6 +30,10 @@ import pytest import torch +from deepmd.dpmodel.utils.dist_check import ( + compute_min_pair_margin_batch, + pair_half_thresholds, +) from deepmd.pt.optimizer import ( HybridMuonOptimizer, ) @@ -46,6 +50,12 @@ from deepmd.utils.compat import ( update_deepmd_input, ) +from deepmd.utils.data import ( + DataRequirementItem, +) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from ..common.stat_file import ( assert_energy_stat_cache_round_trip, @@ -618,6 +628,69 @@ def test_training_loop_compiled(self) -> None: config = normalize(config) self._run_training(config) + @REQUIRES_SUPPORTED_COMPILE + def test_training_loop_compiled_bridged(self) -> None: + """Run compiled training of a ZBL-bridged DPA4C composition. + + The composition has no single fitting net, so the buffers the compiled + graph reads per call belong to its atomic model alone. + """ + config = _make_config(self.data_dir, numb_steps=5) + config["model"] = { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa4c", + "rcut": 4.0, + "channels": 8, + "lmax": 2, + "n_radial": 8, + "seed": 1, + }, + "fitting_net": {"neuron": [16, 16], "seed": 1}, + "bridging_method": "zbl", + "bridging_r_inner": 0.5, + "bridging_r_outer": 0.8, + "data_stat_nbatch": 1, + } + config["training"]["enable_compile"] = True + config = update_deepmd_input(config, warning=False) + config = normalize(config) + self._run_training(config) + + @REQUIRES_SUPPORTED_COMPILE + def test_training_loop_compiled_zbl(self) -> None: + """Run compiled training of a composition on the dense route. + + The tabulated pair potential has no graph lower, so the composition + trains through the dense compiled forward; it has no single fitting + net, so the buffers that forward reads per call belong to its atomic + model alone. + """ + config = _make_config(self.data_dir, numb_steps=5) + config["model"] = { + # The table carries three element columns. + "type_map": ["O", "H", "B"], + "use_srtab": os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "pt", + "model", + "water", + "data", + "zbl_tab_potential", + "H2O_tab_potential.txt", + ), + "smin_alpha": 0.1, + "sw_rmin": 0.2, + "sw_rmax": 1.0, + "descriptor": copy.deepcopy(_DESCRIPTOR_DPA1_NO_ATTN), + "fitting_net": copy.deepcopy(config["model"]["fitting_net"]), + "data_stat_nbatch": 1, + } + config["training"]["enable_compile"] = True + config = update_deepmd_input(config, warning=False) + config = normalize(config) + self._run_training(config) + @REQUIRES_SUPPORTED_COMPILE def test_training_loop_compiled_silu(self) -> None: """Run compiled training with silu activation.""" @@ -1008,6 +1081,146 @@ def test_batch_shapes(self) -> None: shutil.rmtree(tmpdir, ignore_errors=True) +@pytest.mark.parametrize("world_size", [1, 2, 4]) +@pytest.mark.parametrize("has_valid_frame", [False, True]) +def test_replicated_numpy_filter_reads_a_complete_local_pass( + tmp_path: Path, + world_size: int, + has_valid_frame: bool, +) -> None: + """A replicated NumPy reader scans both frames independently of rank count.""" + set_dir = tmp_path / "set.000" + set_dir.mkdir() + np.savetxt(tmp_path / "type.raw", [0, 1], fmt="%d") + np.save( + set_dir / "coord.npy", + np.array( + [ + [0.0, 0.0, 0.0, distance, 0.0, 0.0] + for distance in (0.1, 2.0 if has_valid_frame else 0.2) + ] + ), + ) + np.save(set_dir / "box.npy", np.tile((6.0 * np.eye(3)).reshape(1, 9), (2, 1))) + data = DeepmdDataSystem( + [str(tmp_path)], batch_size=1, test_size=1, type_map=["O", "H"] + ) + data.add_data_requirements( + [DataRequirementItem("pair_margin", 1, default=0.5, source_policy="derived")] + ) + trainer = training_module.Trainer.__new__(training_module.Trainer) + trainer.multi_task = False + trainer.world_size = world_size + trainer.training_data_by_task = {"Default": data} + trainer.min_pair_dist_by_task = {"Default": 0.5} + with ( + patch("deepmd.utils.data.dp_random.shuffle"), + patch.object(data, "get_batch", wraps=data.get_batch) as get_batch, + ): + if has_valid_frame: + _, labels = trainer.get_data(is_train=True) + assert float(labels["pair_margin"].reshape(-1)[0]) >= 1.0 + else: + with pytest.raises(RuntimeError, match="2 consecutive batches"): + trainer.get_data(is_train=True) + assert get_batch.call_count == 2 + + +class TestMinPairDistFilter(unittest.TestCase): + """`training_data.min_pair_dist` drops too-close frames from training.""" + + @classmethod + def setUpClass(cls) -> None: + data_dir = os.path.join(EXAMPLE_DIR, "data") + if not os.path.isdir(data_dir): + raise unittest.SkipTest(f"Example data not found: {data_dir}") + cls.data_dir = data_dir + + def _trainer(self, min_pair_dist: float) -> "training_module.Trainer": + config = _make_config(self.data_dir, numb_steps=5) + config["training"]["training_data"]["batch_size"] = 8 + config["training"]["training_data"]["min_pair_dist"] = min_pair_dist + config = normalize(update_deepmd_input(config, warning=False)) + return get_trainer(config) + + def _run(self, body: Callable[[], None]) -> None: + tmpdir = tempfile.mkdtemp(prefix="pt_expt_min_pair_dist_") + old_cwd = os.getcwd() + os.chdir(tmpdir) + try: + body() + finally: + os.chdir(old_cwd) + shutil.rmtree(tmpdir, ignore_errors=True) + + def test_frames_below_the_threshold_are_dropped(self) -> None: + def body() -> None: + # A frame that clears its window carries a margin it reaches + # rather than its minimum, so the distances that set a splitting + # threshold are measured on the geometry itself. + probe = self._trainer(1.0e-3) + input_dict, label_dict = probe.get_data(is_train=True) + self.assertEqual(label_dict["pair_margin"].numel(), 8) + distances = compute_min_pair_margin_batch( + input_dict["coord"].detach().cpu().numpy(), + input_dict["box"].detach().cpu().numpy(), + input_dict["atype"].detach().cpu().numpy(), + pair_half_thresholds(1.0, "absolute", ntypes=8), + ) + threshold = float(np.median(distances)) + self.assertTrue((distances < threshold).any()) + + trainer = self._trainer(threshold) + sizes = [] + for _ in range(6): + input_dict, label_dict = trainer.get_data(is_train=True) + kept = label_dict["pair_margin"].reshape(-1) + self.assertTrue((kept >= 1.0).all()) + self.assertEqual(input_dict["coord"].shape[0], kept.numel()) + self.assertEqual(label_dict["energy"].shape[0], kept.numel()) + sizes.append(kept.numel()) + self.assertTrue(all(size >= 1 for size in sizes)) + self.assertLess(min(sizes), 8) + # Validation batches are never filtered. + _, valid_label = trainer.get_data(is_train=False) + self.assertNotIn("pair_margin", valid_label) + + self._run(body) + + def test_a_dataset_without_valid_frames_is_reported(self) -> None: + def body() -> None: + # The statistics see the filtered frames, so they fail first. + with self.assertRaisesRegex(RuntimeError, "beyond the filter radius"): + self._trainer(100.0) + # Widening the registered window past every pair empties every + # batch, which the training loop reports after one epoch of retries. + trainer = self._trainer(1.0e-3) + trainer.training_data_by_task[ + training_module.DEFAULT_TASK_KEY + ].add_data_requirements( + [ + DataRequirementItem( + "pair_margin", + ndof=1, + default=100.0, + source_policy="derived", + ) + ] + ) + with self.assertRaisesRegex(RuntimeError, "consecutive batches"): + trainer.get_data(is_train=True) + + self._run(body) + + def test_training_runs_on_the_filtered_batches(self) -> None: + def body() -> None: + trainer = self._trainer(0.9) + trainer.run() + self.assertTrue(os.path.exists("lcurve.out")) + + self._run(body) + + class TestAdditionalDataRequirement(unittest.TestCase): """Cover both branches of fparam handling in get_additional_data_requirement. diff --git a/source/tests/pt_expt/test_training_ddp.py b/source/tests/pt_expt/test_training_ddp.py index 3afbda62e3..f160de446c 100644 --- a/source/tests/pt_expt/test_training_ddp.py +++ b/source/tests/pt_expt/test_training_ddp.py @@ -361,6 +361,55 @@ def _worker_single_task_train(rank, world_size, port, data_dir, result_dict): dist.destroy_process_group() +def _worker_filtered_train(rank, world_size, port, data_dir, result_dict): + """Worker: DDP training whose ranks drop different too-close frames.""" + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + dist.init_process_group(backend=_DDP_BACKEND, rank=rank, world_size=world_size) + try: + tmpdir = tempfile.mkdtemp(prefix=f"ddp_filter_rank{rank}_") + old_cwd = os.getcwd() + os.chdir(tmpdir) + try: + config = _make_config(data_dir, numb_steps=4) + # About a quarter of the water frames keep every pair 0.92 Å + # apart, so most batches shrink and some are replaced outright. + config["training"]["training_data"]["batch_size"] = 2 + config["training"]["training_data"]["min_pair_dist"] = 0.92 + config["training"]["seed"] = 10 + rank + config = update_deepmd_input(config, warning=False) + config = normalize(config) + trainer = get_trainer(config) + sizes = [] + fetch = trainer.get_data + + def counting_fetch(*args, **kwargs): + batch = fetch(*args, **kwargs) + if kwargs.get("is_train", True): + kept = batch[1]["pair_margin"].reshape(-1) + # The field is a margin against the frame's own pair + # thresholds, so every surviving frame reaches one; a + # frame that slipped through unfiltered would not. + assert bool((kept >= 1.0).all()) + sizes.append(int(kept.numel())) + return batch + + trainer.get_data = counting_fetch + trainer.run() + result_dict[rank] = { + "sizes": sizes, + "weights": { + name: p.detach().cpu().clone() + for name, p in trainer._unwrapped.named_parameters() + }, + } + finally: + os.chdir(old_cwd) + shutil.rmtree(tmpdir, ignore_errors=True) + finally: + dist.destroy_process_group() + + def _worker_multitask_train(rank, world_size, port, data_dir, result_dict): """Worker: run multi-task DDP training.""" os.environ["MASTER_ADDR"] = "127.0.0.1" @@ -940,6 +989,43 @@ def test_ddp_single_task_trains(self) -> None: ) +class TestDDPMinPairDistFilter(unittest.TestCase): + """`min_pair_dist` filters per rank without breaking the DDP lockstep.""" + + @classmethod + def setUpClass(cls) -> None: + data_dir = os.path.join(EXAMPLE_DIR, "data") + if not os.path.isdir(data_dir): + raise unittest.SkipTest(f"Example data not found: {data_dir}") + cls.data_dir = os.path.join(data_dir, "data_0") + + def test_ranks_stay_in_lockstep(self) -> None: + port = _find_free_port() + result_dict = mp.Manager().dict() + mp.spawn( + _worker_filtered_train, + args=(2, port, self.data_dir, result_dict), + nprocs=2, + join=True, + ) + results = dict(result_dict) + # Every rank contributes a non-empty batch to every step, and the + # filter is active: at least one batch lost a frame. + for rank in (0, 1): + self.assertGreaterEqual(len(results[rank]["sizes"]), 4) + self.assertTrue(all(size >= 1 for size in results[rank]["sizes"])) + self.assertLess( + min(results[0]["sizes"] + results[1]["sizes"]), + 2, + ) + for name in results[0]["weights"]: + torch.testing.assert_close( + results[0]["weights"][name], + results[1]["weights"][name], + msg=f"Weights differ across ranks: {name}", + ) + + class TestDDPMultiTaskTrain(unittest.TestCase): """Smoke test: multi-task DDP training with 2 ranks.""" diff --git a/source/tests/tf2/test_dpa4.py b/source/tests/tf2/test_dpa4.py index 6c5e2bf686..c2e2cf13bc 100644 --- a/source/tests/tf2/test_dpa4.py +++ b/source/tests/tf2/test_dpa4.py @@ -285,7 +285,7 @@ def run(n_node, atype, edge_index, edge_vec, edge_mask): edge_vec=wrap_tensor(edge_vec), edge_mask=wrap_tensor(edge_mask), ) - descrpt, _ = descriptor.call_graph(graph, wrap_tensor(atype)) + descrpt, _, _ = descriptor.call_graph(graph, wrap_tensor(atype)) return to_tf_tensor(descrpt) args = (