diff --git a/deepmd/dpmodel/atomic_model/__init__.py b/deepmd/dpmodel/atomic_model/__init__.py index c55eead5f3..be6b2ea8ac 100644 --- a/deepmd/dpmodel/atomic_model/__init__.py +++ b/deepmd/dpmodel/atomic_model/__init__.py @@ -48,6 +48,9 @@ from .unimol_atomic_model import ( DPUniMolAtomicModel, ) +from .unimol_dpa_atomic_model import ( + DPUniMolDPAAtomicModel, +) __all__ = [ "BaseAtomicModel", @@ -58,6 +61,7 @@ "DPPolarAtomicModel", "DPPropertyAtomicModel", "DPUniMolAtomicModel", + "DPUniMolDPAAtomicModel", "DPZBLLinearEnergyAtomicModel", "LinearEnergyAtomicModel", "PairTabAtomicModel", diff --git a/deepmd/dpmodel/atomic_model/unimol_dpa_atomic_model.py b/deepmd/dpmodel/atomic_model/unimol_dpa_atomic_model.py new file mode 100644 index 0000000000..9d08da8de5 --- /dev/null +++ b/deepmd/dpmodel/atomic_model/unimol_dpa_atomic_model.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Atomic model for Uni-Mol pretraining on a DPA backbone.""" + +from typing import ( + Any, +) + +from deepmd.dpmodel.array_api import ( + Array, +) +from deepmd.dpmodel.fitting.unimol_dpa_pretrain import ( + UniMolDPAPretrainFitting, +) + +from .dp_atomic_model import ( + DPAtomicModel, +) + + +class DPUniMolDPAAtomicModel(DPAtomicModel): + r"""Uni-Mol's objective wired onto a DPA backbone. + + The standard path hands a descriptor's five-tuple to a fitting. Two of the + things these heads read are not in it: the equivariant state the coordinate + head projects, which the descriptor otherwise discards, and the neighbour + list the distance head uses to decide which pairs it covers. This model + overrides one method to fetch both and hand them over. Nothing else about + the atomic model changes. + """ + + def __init__( + self, descriptor: Any, fitting: Any, type_map: list[str], **kwargs: Any + ) -> None: + if not hasattr(descriptor, "call_with_latent"): + raise TypeError( + "the unimol DPA objective needs a backbone that exposes its " + "equivariant state, which is what the coordinate head reads; " + f"{type(descriptor).__name__} does not have call_with_latent" + ) + if not isinstance(fitting, UniMolDPAPretrainFitting): + raise TypeError( + "DPUniMolDPAAtomicModel needs the unimol_dpa_pretrain fitting" + ) + super().__init__(descriptor, fitting, type_map, **kwargs) + + def forward_atomic( + self, + extended_coord: Array, + extended_atype: Array, + nlist: Array, + mapping: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, + comm_dict: dict | None = None, + charge_spin: Array | None = None, + ) -> dict[str, Array]: + """Run the backbone and its heads. + + Parameters + ---------- + extended_coord + nf x (nall x 3) coordinates. + extended_atype + nf x nall element types, already clamped to be non-negative. + nlist + nf x nloc x nnei neighbour list, which also decides which pairs the + distance head covers. + mapping + Extended-to-local mapping, passed through to the backbone. + fparam, aparam, comm_dict + Unused by this model. + charge_spin + nf x 2 frame charge and spin, passed through to the backbone, + which fills in its ``default_chg_spin`` when this is None. + + Returns + ------- + dict + The head outputs, plus the mask saying which pairs are covered. + """ + del fparam, aparam, comm_dict + # Every path reaches the objective here, including the lower one that + # evaluation and export drive, so this is where a periodic frame has to + # be refused: ghosts are the shape a cell takes by the time it arrives. + # The coverage mask keeps only local neighbours and the distance target + # is a plain coordinate difference with no minimum image, so a cell + # would train against wrong labels rather than fail. + nall = extended_atype.shape[1] + nloc = nlist.shape[1] + if nall > nloc: + raise ValueError( + f"the unimol objective needs every atom to be local, but this " + f"frame carries {nall - nloc} ghost atom(s), which is what a " + "periodic cell becomes here. Its distance target has no " + "minimum-image convention and its coverage keeps only local " + "neighbours, so a cell would quietly mistrain rather than fail" + ) + node_ebd, latent = self.descriptor.call_with_latent( + extended_coord, + extended_atype, + nlist, + mapping=mapping, + charge_spin=charge_spin, + ) + return self.fitting_net.call_atoms(node_ebd, latent, nlist, extended_coord) + + def supports_graph_export(self) -> bool: + """No graph entry: these heads read the neighbour list directly. + + The base class advertises one, and the export path believes it, so + without this a graph export would reach the fitting and fail on a + missing ``call_graph`` rather than be refused up front. + """ + return False + + def uses_graph_lower(self) -> bool: + """See :meth:`supports_graph_export`.""" + return False + + def apply_out_stat(self, ret: dict[str, Array], atype: Array) -> dict[str, Array]: + """Return the head outputs untouched. + + Self-supervised targets carry no per-element bias to add back: the + element head predicts a distribution, and the coordinate and distance + heads predict geometry the data already fixes. + """ + del atype + return ret diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index 59c027bbf9..04fcf4ea01 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -1284,6 +1284,7 @@ def call( force_embedding: Array | None = None, charge_spin: Array | None = None, spin: Array | None = None, + return_latent: bool = False, ) -> tuple[ Array, Array | None, @@ -1319,6 +1320,9 @@ def call( initial SO(3) backbone state before the interaction blocks. charge_spin Frame-level charge and spin conditions with shape (nf, 2). + return_latent + Return ``(descriptor, latent)`` instead of the five-tuple. Prefer + :meth:`call_with_latent`, which is this flag with a name. Returns ------- @@ -1381,7 +1385,7 @@ def call( graph, atype_flat = _graph_from_padded_nlist( coord_ext, atype_ext, nlist, mapping ) - x_scalar, _ = self._run_graph( + x_scalar, latent = self._run_graph( graph, atype_flat, n_out_nodes=nf * nloc, @@ -1392,12 +1396,72 @@ 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)) - return ( - descriptor, - None, - None, - None, - None, + if not return_latent: + return ( + descriptor, + None, + None, + None, + None, + ) + # Cast the latent the way ``_run_graph`` already casts the scalar, so + # that everything leaving a public entry is in the global precision. + return descriptor, xp.astype(latent, get_xp_precision(xp, "global")) + + def call_with_latent( + self, + coord_ext: Array, + atype_ext: Array, + nlist: Array, + mapping: Array | None = None, + comm_dict: dict[str, Array] | None = None, + fparam: Array | None = None, + force_embedding: Array | None = None, + charge_spin: Array | None = None, + spin: Array | None = None, + ) -> tuple[Array, Array]: + """Compute the descriptor and the equivariant state it was read out of. + + ``call`` returns the l=0 read-out only, which is all a potential energy + surface needs. A head that predicts a per-atom vector needs the l=1 part + of the backbone state instead, and that state is otherwise discarded. + The pt backend exposes the same thing as ``forward_with_edges``. + + Arguments are those of :meth:`call`. + + Returns + ------- + descriptor + Descriptor with shape (nf, nloc, channels), exactly what + :meth:`call` returns in its first slot. + latent + The final equivariant state, with shape + ``(nf * nloc, node_readout_dim, 1, channels)`` where + ``node_readout_dim = (node_readout_lmax + 1) ** 2``. Rows ``1:4`` of + the second axis are the l=1 block, in a packed real-harmonic basis + rather than Cartesian order; see + :func:`deepmd.dpmodel.fitting.unimol_dpa_heads.l1_to_cartesian` for + the mapping that rotates with the input, which was established by + measurement rather than by assuming a convention. + + Note that on the torch backend this state is not bit-reproducible + between identical calls when ``use_env_seed`` is on and the + precision is single: measured at 7e-09, against exactly zero for the + scalar read-out, and gone in double precision or with + ``use_env_seed`` off. The scalar output :meth:`call` returns is + unaffected, so this surfaces only through this accessor. + """ + return self.call( + coord_ext, + atype_ext, + nlist, + mapping=mapping, + comm_dict=comm_dict, + fparam=fparam, + force_embedding=force_embedding, + charge_spin=charge_spin, + spin=spin, + return_latent=True, ) def call_graph( diff --git a/deepmd/dpmodel/fitting/unimol_dpa_heads.py b/deepmd/dpmodel/fitting/unimol_dpa_heads.py new file mode 100644 index 0000000000..896df1e044 --- /dev/null +++ b/deepmd/dpmodel/fitting/unimol_dpa_heads.py @@ -0,0 +1,321 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Uni-Mol pretraining heads that read a DPA backbone. + +Uni-Mol's own heads read a pair representation, which its transformer carries +and a DPA backbone does not: DPA4 attends per edge with a scatter softmax and +returns no pair axis at all. These heads read what DPA4 does expose instead -- +the l=0 descriptor, and the equivariant state behind it. +""" + +import array_api_compat + +from deepmd.dpmodel.array_api import ( + Array, +) +from deepmd.dpmodel.common import ( + NativeOP, +) +from deepmd.dpmodel.descriptor.dpa4_nn.so3 import ( + SO3Linear, +) +from deepmd.dpmodel.utils.network import ( + LayerNorm, + NativeLayer, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) + + +def l1_to_cartesian(l1: Array) -> Array: + r"""Turn an l=1 block in SeZM's packed basis into a Cartesian vector. + + The three rows of the l=1 block are not :math:`(x, y, z)`. Reading them as + such gives a quantity that does not rotate with the input, which is the + whole point of taking them. The mapping below is the one that does: + + .. math:: + + (x, y, z) = (-c_2,\; c_0,\; c_1) + + It was established by rotating the input by a random :math:`R \in SO(3)` and + asking which of the 48 signed permutations of the three rows satisfies + :math:`v(Rx) = R\,v(x)`. Two do, to 1.7e-14 in float64 -- this one and its + negation, as must be the case, since negating a map that commutes with the + rotation leaves it commuting. The next best is wrong by 0.2, so the pair is + unambiguous, and which of the two is used does not matter: the projection + feeding this is learned and absorbs an overall sign. + + Parameters + ---------- + l1 : Array + The l=1 rows, shape ``(..., 3)``. + + Returns + ------- + Array + Cartesian vectors, shape ``(..., 3)``. + """ + xp = array_api_compat.array_namespace(l1) + return xp.stack([-l1[..., 2], l1[..., 0], l1[..., 1]], axis=-1) + + +class EquivariantCoordHead(NativeOP): + r"""Predict a per-atom coordinate update from a DPA backbone's l=1 state. + + Uni-Mol predicts the update from its pair channel, weighting each + :math:`x_j - x_i` by a learned scalar. A DPA backbone has no pair channel, + so the update is read from the l=1 part of the equivariant state instead, + which is a per-atom vector by construction. The projection is degree-wise + and shares its weights across the three :math:`m` components, so the result + rotates with the input rather than merely being three numbers. + + Parameters + ---------- + lmax : int + Degree of the state being read; ``node_readout_lmax`` of the descriptor. + channels : int + Channel width of the state. + precision : str + Floating-point precision of the parameters. + seed : int, optional + Random seed for initialization. + """ + + def __init__( + self, + lmax: int, + channels: int, + precision: str = "float64", + seed: int | list[int] | None = None, + ) -> None: + if lmax < 1: + raise ValueError( + f"a coordinate head needs the l=1 block, so the backbone must " + f"read out at least degree 1; got lmax={lmax}" + ) + self.lmax = int(lmax) + self.channels = int(channels) + self.precision = precision + self.proj = SO3Linear( + lmax=self.lmax, + in_channels=self.channels, + out_channels=1, + n_focus=1, + precision=precision, + seed=seed, + ) + + def call(self, latent: Array) -> Array: + """Read one Cartesian vector per atom. + + Parameters + ---------- + latent : Array + The descriptor's equivariant state, shape + ``(nf * nloc, (lmax + 1) ** 2, 1, channels)``. + + Returns + ------- + Array + Per-atom update, shape ``(nf * nloc, 3)``. + """ + projected = self.proj(latent) + return l1_to_cartesian(projected[:, 1:4, 0, 0]) + + def serialize(self) -> dict: + """Serialize the head.""" + return { + "@class": "EquivariantCoordHead", + "@version": 1, + "lmax": self.lmax, + "channels": self.channels, + "precision": self.precision, + "proj": self.proj.serialize(), + } + + @classmethod + def deserialize(cls, data: dict) -> "EquivariantCoordHead": + """Deserialize the head.""" + data = data.copy() + data.pop("@class", None) + data.pop("@version", None) + proj = data.pop("proj") + obj = cls(**data) + obj.proj = SO3Linear.deserialize(proj) + return obj + + +class PairDistanceHead(NativeOP): + r"""Predict clean pairwise distances without a pair representation. + + Uni-Mol reads this from its pair channel, which carries a vector per ordered + atom pair. DPA4 has no such channel, so the pair is described by its two + endpoints instead -- representations that have already exchanged information + through the backbone's message passing. The combination is symmetric in the + two endpoints, because a distance is: + + .. math:: + + p_{ij} = \big[\, h_i + h_j \;\|\; h_i \odot h_j \,\big] + + Which pairs are covered is a choice, and it is the one place this objective + departs from Uni-Mol by construction: + + ``neighbour`` + Only pairs inside the backbone's neighbour list. The mask reuses the + structure the backbone already built, but the dense pair features and + MLP are still evaluated for every ``nloc x nloc`` pair, so this setting + has O(nloc^2) head compute and only reduces the scored entries. On + drug-like molecules a 6 Angstrom cut-off holds about half of all pairs, + so the objective sees less than Uni-Mol's. + ``all_pairs`` + Every ordered pair, the diagonal included, which is Uni-Mol's own + coverage and what to use to reproduce its training. Costs O(nloc^2) and + ignores the neighbour list, so it does not share the backbone's + locality. + + Parameters + ---------- + dim_descrpt : int + Width of the per-atom representation. + hidden : int + Width of the head's hidden layer. + coverage : str + ``neighbour`` or ``all_pairs``. + activation_function : str + Activation of the head. + precision : str + Floating-point precision of the parameters. + seed : int, optional + Random seed for initialization. + """ + + COVERAGES = ("neighbour", "all_pairs") + + def __init__( + self, + dim_descrpt: int, + hidden: int = 64, + coverage: str = "neighbour", + activation_function: str = "gelu_erf", + precision: str = "float64", + seed: int | list[int] | None = None, + ) -> None: + if coverage not in self.COVERAGES: + raise ValueError( + f"unknown coverage {coverage!r}; it must be one of " + f"{', '.join(map(repr, self.COVERAGES))}" + ) + self.dim_descrpt = int(dim_descrpt) + self.hidden = int(hidden) + self.coverage = coverage + self.activation_function = activation_function + self.precision = precision + self.dense = NativeLayer( + 2 * self.dim_descrpt, + self.hidden, + activation_function=activation_function, + precision=precision, + seed=child_seed(seed, 0), + ) + self.layer_norm = LayerNorm( + self.hidden, precision=precision, seed=child_seed(seed, 1) + ) + self.out_proj = NativeLayer( + self.hidden, 1, bias=True, precision=precision, seed=child_seed(seed, 2) + ) + + def call(self, node_ebd: Array, nlist: Array) -> tuple[Array, Array]: + """Predict a distance per covered pair. + + Parameters + ---------- + node_ebd : Array + Per-atom representation, shape ``(nf, nloc, dim_descrpt)``. + nlist : Array + Neighbour list, shape ``(nf, nloc, nnei)``. Negative entries are + padding. Read only when the coverage is ``neighbour``. + + Returns + ------- + pair_dist : Array + Predicted distances, shape ``(nf, nloc, nloc)``. + pair_mask : Array + 1 where the pair is covered, else 0, same shape. The objective + averages over these entries only. What "covered" means is the + coverage setting: a neighbour list never holds an atom as its own + neighbour, so ``neighbour`` leaves the diagonal at 0; ``all_pairs`` + reproduces upstream, which scores the zero self-distance too, so + its diagonal is 1. + """ + xp = array_api_compat.array_namespace(node_ebd) + nf, nloc = node_ebd.shape[0], node_ebd.shape[1] + dev = array_api_compat.device(node_ebd) + + summed = node_ebd[:, :, None, :] + node_ebd[:, None, :, :] + product = node_ebd[:, :, None, :] * node_ebd[:, None, :, :] + pair = xp.concat([summed, product], axis=-1) + predicted = self.out_proj(self.layer_norm(self.dense(pair))) + predicted = xp.reshape(predicted, (nf, nloc, nloc)) + + if self.coverage == "neighbour": + # A neighbour list never lists an atom as its own neighbour, so the + # diagonal is excluded by construction rather than by subtracting it. + mask = _neighbour_mask(nlist, nloc, node_ebd.dtype) + else: + # Every pair, the diagonal included: upstream scores the zero + # self-distance too (losses/unimol.py:159-184 takes every + # non-padding column), and those entries carry the most extreme + # normalised label in the set. Dropping them would make this + # setting something other than the coverage it claims to reproduce. + mask = xp.ones((nf, nloc, nloc), dtype=node_ebd.dtype, device=dev) + return predicted, mask + + def serialize(self) -> dict: + """Serialize the head.""" + return { + "@class": "PairDistanceHead", + "@version": 1, + "dim_descrpt": self.dim_descrpt, + "hidden": self.hidden, + "coverage": self.coverage, + "activation_function": self.activation_function, + "precision": self.precision, + "dense": self.dense.serialize(), + "layer_norm": self.layer_norm.serialize(), + "out_proj": self.out_proj.serialize(), + } + + @classmethod + def deserialize(cls, data: dict) -> "PairDistanceHead": + """Deserialize the head.""" + data = data.copy() + data.pop("@class", None) + data.pop("@version", None) + parts = {k: data.pop(k) for k in ("dense", "layer_norm", "out_proj")} + obj = cls(**data) + obj.dense = NativeLayer.deserialize(parts["dense"]) + obj.layer_norm = LayerNorm.deserialize(parts["layer_norm"]) + obj.out_proj = NativeLayer.deserialize(parts["out_proj"]) + return obj + + +def _neighbour_mask(nlist: Array, nloc: int, dtype) -> Array: # noqa: ANN001 + """Scatter a padded neighbour list into a dense ``(nf, nloc, nloc)`` mask. + + One neighbour slot at a time, rather than one comparison across all of them: + the latter reads better but allocates ``nf * nloc * nnei * nloc`` values, + which at a realistic ``sel`` and batch size is gigabytes per step for a + result of a few megabytes. The loop is over a static dimension. + """ + xp = array_api_compat.array_namespace(nlist) + dev = array_api_compat.device(nlist) + columns = xp.arange(nloc, dtype=nlist.dtype, device=dev)[None, None, :] + mask = xp.zeros((nlist.shape[0], nlist.shape[1], nloc), dtype=dtype, device=dev) + for k in range(nlist.shape[2]): + # Padding is negative and a neighbour beyond the local atoms is a ghost; + # neither names a pair this objective can score, and neither matches any + # column, so both contribute nothing. + mask = xp.maximum(mask, xp.astype(nlist[:, :, k : k + 1] == columns, dtype)) + return mask diff --git a/deepmd/dpmodel/fitting/unimol_dpa_pretrain.py b/deepmd/dpmodel/fitting/unimol_dpa_pretrain.py new file mode 100644 index 0000000000..a701ffddfd --- /dev/null +++ b/deepmd/dpmodel/fitting/unimol_dpa_pretrain.py @@ -0,0 +1,361 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Uni-Mol's pretraining objective on a DPA backbone. + +The three heads are the same three Uni-Mol trains: the element of a corrupted +atom, the clean coordinates, and the clean pairwise distances. What changes is +what they read. Uni-Mol's transformer carries a pair representation and wraps +each molecule in two virtual tokens; a DPA backbone carries neither, so the +coordinate head reads the equivariant state and the distance head reads pairs of +atom representations. The element head is unchanged in kind -- it reads the +per-atom representation either way. + +The two regularisers Uni-Mol places on its own encoder, on the node norm and on +the pair-delta norm, have no counterpart here: they constrain quantities that +belong to that transformer. The objective leaves them at zero weight. +""" + +import array_api_compat + +from deepmd.dpmodel.array_api import ( + Array, +) +from deepmd.dpmodel.common import ( + NativeOP, + safe_cast_array, +) +from deepmd.dpmodel.descriptor.unimol_nn import ( + MaskLMHead, +) +from deepmd.dpmodel.fitting.base_fitting import ( + BaseFitting, +) +from deepmd.dpmodel.fitting.unimol_dpa_heads import ( + EquivariantCoordHead, + PairDistanceHead, +) +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + OutputVariableDef, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + + +@BaseFitting.register("unimol_dpa_pretrain") +class UniMolDPAPretrainFitting(NativeOP, BaseFitting): + r"""Uni-Mol's self-supervised heads, reading a DPA backbone. + + Parameters + ---------- + ntypes : int + Number of element types of the model. + dim_descrpt : int + Width of the per-atom representation the backbone produces. + node_readout_lmax : int + Degree the backbone reads out, which fixes the shape of the equivariant + state the coordinate head consumes. Must be at least 1. + n_token : int + Size of the Uni-Mol vocabulary, which is the width of the element head. + max_atoms : int + Largest molecule accepted. The distance head declares a fixed number of + columns, so this fixes it; larger frames cannot be expressed. + dist_hidden : int + Width of the distance head's hidden layer. + dist_coverage : str + Which pairs the distance head covers: ``neighbour`` follows the + backbone's neighbour list, ``all_pairs`` covers every pair as Uni-Mol + does. See :class:`PairDistanceHead`. + activation_function : str + Activation of the heads. + mask_token_head, coord_head, dist_head : bool + Build each head. + precision : str + Floating-point precision of the parameters. + seed : int, optional + Random seed for initialization. + type_map : list of str, optional + Element names, as configured on the model. + """ + + def __init__( + self, + ntypes: int, + dim_descrpt: int, + node_readout_lmax: int, + n_token: int = 31, + max_atoms: int = 256, + dist_hidden: int = 64, + dist_coverage: str = "neighbour", + activation_function: str = "gelu_erf", + mask_token_head: bool = True, + coord_head: bool = True, + dist_head: bool = True, + precision: str = "float64", + seed: int | list[int] | None = None, + type_map: list[str] | None = None, + **kwargs: float | str | bool, + ) -> None: + del kwargs + self.type_map = list(type_map) if type_map is not None else None + self.ntypes = ntypes + self.dim_descrpt = dim_descrpt + self.node_readout_lmax = node_readout_lmax + self.n_token = n_token + self.max_atoms = max_atoms + self.dist_hidden = dist_hidden + self.dist_coverage = dist_coverage + self.activation_function = activation_function + self.precision = precision + self.lm_head = ( + MaskLMHead( + dim_descrpt, + n_token, + activation_function, + precision, + child_seed(seed, 0), + ) + if mask_token_head + else None + ) + self.coord_head = ( + EquivariantCoordHead( + lmax=node_readout_lmax, + channels=dim_descrpt, + precision=precision, + seed=child_seed(seed, 1), + ) + if coord_head + else None + ) + self.dist_head = ( + PairDistanceHead( + dim_descrpt=dim_descrpt, + hidden=dist_hidden, + coverage=dist_coverage, + activation_function=activation_function, + precision=precision, + seed=child_seed(seed, 2), + ) + if dist_head + else None + ) + + def get_type_map(self) -> list[str]: + """Element names, as configured on the model.""" + return self.type_map if self.type_map is not None else [] + + def change_type_map( + self, type_map: list[str], model_with_new_type_stat: object = None + ) -> None: + """Adopt a new element list. + + The element head is indexed by Uni-Mol token rather than by deepmd type, + so nothing but the recorded names changes. + """ + del model_with_new_type_stat + self.type_map = list(type_map) + self.ntypes = len(type_map) + + def compute_input_stats(self, merged, stat_file_path=None, **kwargs) -> None: # noqa: ANN001, ANN003 + """No input statistics: the heads read a learned representation.""" + + def get_dim_fparam(self) -> int: + """No frame parameters: the objective reads structure only.""" + return 0 + + def get_dim_aparam(self) -> int: + """No atomic parameters.""" + return 0 + + def has_default_fparam(self) -> bool: + """There are no frame parameters, so there is no default either.""" + return False + + def get_default_fparam(self): # noqa: ANN201 + """There are no frame parameters.""" + return None + + def get_sel_type(self) -> list[int]: + """Every element takes part in the objective.""" + return [] + + def reinit_exclude(self, exclude_types: list[int] | None = None) -> None: + """Type exclusion is meaningless here: every atom is predicted.""" + if exclude_types: + raise NotImplementedError( + "unimol_dpa_pretrain predicts every atom and does not support " + "excluded types" + ) + + def set_case_embd(self, case_idx: int) -> None: + """Case embeddings are a multi-task feature these heads do not use.""" + raise NotImplementedError( + "unimol_dpa_pretrain does not support case embeddings" + ) + + def output_def(self) -> FittingOutputDef: + """Declare the head outputs. + + All are per-atom, none reduces to a frame total, and none is + differentiated with respect to coordinates or cell: this objective + denoises a structure rather than predicting a potential energy surface. + """ + variables = [] + if self.lm_head is not None: + variables.append( + OutputVariableDef( + "token_logits", + [self.n_token], + reducible=False, + r_differentiable=False, + c_differentiable=False, + ) + ) + if self.coord_head is not None: + variables.append( + OutputVariableDef( + "coord_update", + [3], + reducible=False, + r_differentiable=False, + c_differentiable=False, + ) + ) + if self.dist_head is not None: + # A fixed column count, because the output machinery indexes this + # definition for every key the heads return. The frame's own atom + # count is shorter, so the rows are padded and the mask says so. + for name in ("pair_dist", "pair_mask"): + variables.append( + OutputVariableDef( + name, + [self.max_atoms], + reducible=False, + r_differentiable=False, + c_differentiable=False, + ) + ) + return FittingOutputDef(variables) + + def call_atoms( + self, node_ebd: Array, latent: Array, nlist: Array, coord: Array + ) -> dict: + """Run the heads on one frame's worth of backbone output. + + Parameters + ---------- + node_ebd : Array + Per-atom representation, shape ``(nf, nloc, dim_descrpt)``. + latent : Array + The backbone's equivariant state, shape + ``(nf * nloc, (lmax + 1) ** 2, 1, dim_descrpt)``. + nlist : Array + Neighbour list, shape ``(nf, nloc, nnei)``. + coord : Array + The input coordinates, shape ``(nf, nall, 3)`` or flattened. The + coordinate head predicts a displacement; adding these local input + coordinates gives the clean-coordinate prediction compared by the + loss, matching Uni-Mol's ``coord + update`` head. + + Returns + ------- + dict + ``token_logits`` and ``coord_update`` per atom; ``pair_dist`` and + the ``pair_mask`` saying which pairs it covers. + """ + node_ebd = safe_cast_array(node_ebd, "global", self.precision) + latent = safe_cast_array(latent, "global", self.precision) + coord = safe_cast_array(coord, "global", self.precision) + xp = array_api_compat.array_namespace(node_ebd) + nf, nloc = node_ebd.shape[0], node_ebd.shape[1] + local_coord = xp.reshape(coord, (nf, -1, 3))[:, :nloc, :] + out = {} + if self.lm_head is not None: + out["token_logits"] = self.lm_head(node_ebd) + if self.coord_head is not None: + displacement = xp.reshape(self.coord_head(latent), (nf, nloc, 3)) + out["coord_update"] = local_coord + displacement + if self.dist_head is not None: + if nloc > self.max_atoms: + raise ValueError( + f"a frame of {nloc} atoms exceeds max_atoms={self.max_atoms}; " + "the distance head declares a fixed width, so larger frames " + "cannot be expressed. Convert the data with a matching " + "max_atoms, or raise it here" + ) + dist, mask = self.dist_head(node_ebd, nlist) + pad = self.max_atoms - nloc + if pad: + shape = (nf, nloc, pad) + dev = array_api_compat.device(node_ebd) + zeros = xp.zeros(shape, dtype=dist.dtype, device=dev) + dist = xp.concat([dist, zeros], axis=-1) + mask = xp.concat( + [mask, xp.zeros(shape, dtype=mask.dtype, device=dev)], axis=-1 + ) + out["pair_dist"] = dist + out["pair_mask"] = mask + return { + kk: safe_cast_array(vv, self.precision, "global") for kk, vv in out.items() + } + + def call(self, descriptor: Array, atype: Array, **kwargs) -> dict[str, Array]: # noqa: ANN003 + """Not reachable: the heads need the equivariant state and the nlist. + + The pretraining path goes through :meth:`call_atoms`, driven by the + atomic model, because neither the equivariant state nor the neighbour + list fits the descriptor's five-tuple. + """ + raise NotImplementedError( + "unimol_dpa_pretrain reads the backbone's equivariant state and its " + "neighbour list; it is driven through call_atoms by the atomic model" + ) + + def serialize(self) -> dict: + """Serialize the fitting.""" + return { + "@class": "Fitting", + "type": "unimol_dpa_pretrain", + "@version": 1, + "ntypes": self.ntypes, + "type_map": self.type_map, + "dim_descrpt": self.dim_descrpt, + "node_readout_lmax": self.node_readout_lmax, + "n_token": self.n_token, + "max_atoms": self.max_atoms, + "dist_hidden": self.dist_hidden, + "dist_coverage": self.dist_coverage, + "activation_function": self.activation_function, + "precision": self.precision, + "mask_token_head": self.lm_head is not None, + "coord_head": self.coord_head is not None, + "dist_head": self.dist_head is not None, + "lm_head": None if self.lm_head is None else self.lm_head.serialize(), + "coord_head_net": None + if self.coord_head is None + else self.coord_head.serialize(), + "dist_head_net": None + if self.dist_head is None + else self.dist_head.serialize(), + } + + @classmethod + def deserialize(cls, data: dict) -> "UniMolDPAPretrainFitting": + """Deserialize the fitting.""" + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + data.pop("@class", None) + data.pop("type", None) + parts = {k: data.pop(k) for k in ("lm_head", "coord_head_net", "dist_head_net")} + obj = cls(**data) + if parts["lm_head"] is not None: + obj.lm_head = MaskLMHead.deserialize(parts["lm_head"]) + if parts["coord_head_net"] is not None: + obj.coord_head = EquivariantCoordHead.deserialize(parts["coord_head_net"]) + if parts["dist_head_net"] is not None: + obj.dist_head = PairDistanceHead.deserialize(parts["dist_head_net"]) + return obj diff --git a/deepmd/dpmodel/loss/unimol.py b/deepmd/dpmodel/loss/unimol.py index 08c464ffc6..c0ac050f11 100644 --- a/deepmd/dpmodel/loss/unimol.py +++ b/deepmd/dpmodel/loss/unimol.py @@ -11,9 +11,15 @@ Five terms, with upstream's default weights from its README pretraining recipe: element prediction (1), coordinate denoising (5), distance prediction (10), and the two norm regularisers (0.01 each). The regularisers are produced by the -backbone, so the loss only weights them. +backbone, so the loss only weights them. A DPA backbone has no corresponding +outputs; its configuration sets both weights to zero and the loss rejects a +nonzero value before training starts. """ +from collections.abc import ( + Iterable, +) + import array_api_compat from deepmd.dpmodel.array_api import ( @@ -34,6 +40,19 @@ DIST_MEAN = 6.312581655060595 DIST_STD = 3.3899264663911888 +# Which model output each weighted term reads. A backbone carries the heads its +# fitting was configured with, and Uni-Mol's two norm regularisers describe its +# own transformer, so not every term can be evaluated against every backbone. +# The weights are what the configuration exposes, so a term that cannot be fed +# is reported by the name of the knob that switches it off. +TERM_OUTPUT = { + "masked_token_loss": "token_logits", + "masked_coord_loss": "coord_update", + "masked_dist_loss": "pair_dist", + "x_norm_loss": "x_norm", + "delta_pair_repr_norm_loss": "delta_pair_norm", +} + def _smooth_l1(pred: Array, label: Array, beta: float = 1.0) -> Array: r"""Mean smooth L1, matching ``F.smooth_l1_loss(reduction="mean")``. @@ -80,30 +99,42 @@ def _frame_scalar(value: Array, mask: Array | None) -> Array: return xp.sum(per_atom * weights) / xp.where(total > 0, total, xp.ones_like(total)) -def _token_mask_from_atoms(mask: Array, ncol: int) -> Array: - """Mark the non-padding token columns: BOS, the real atoms, then EOS.""" +def _token_mask_from_atoms(mask: Array, ncol: int, virtual: bool = True) -> Array: + """Mark the non-padding token columns. + + For Uni-Mol those are BOS, the real atoms, then EOS. A DPA backbone wraps + the molecule in no virtual tokens at all, so there the columns are just the + real atoms. + """ xp = array_api_compat.array_namespace(mask) n_real = xp.sum(xp.astype(mask, xp.int64), axis=-1) positions = xp.arange(ncol, dtype=xp.int64, device=array_api_compat.device(mask))[ None, : ] - return xp.astype(positions < (n_real + 2)[:, None], xp.int64) + return xp.astype(positions < (n_real + (2 if virtual else 0))[:, None], xp.int64) -def _clean_distances(coord_target: Array, mask: Array, ncol: int) -> Array: - """Pairwise distances from the clean coordinates, virtual tokens included. +def _clean_distances( + coord_target: Array, mask: Array, ncol: int, virtual: bool = True +) -> Array: + """Pairwise distances from the clean coordinates. Storing this as a label would cost O(natoms^2) per frame, so it is derived - here instead. The two virtual tokens sit at the origin, which is the - centroid of the clean coordinates because the transform centres them. + here instead. With ``virtual``, the two virtual tokens are included and sit + at the origin, which is the centroid of the clean coordinates because the + transform centres them. A DPA backbone has no virtual tokens and passes + ``virtual=False``. """ xp = array_api_compat.array_namespace(coord_target) nf = coord_target.shape[0] real = xp.astype(mask, coord_target.dtype)[..., None] atoms = coord_target * real dev = array_api_compat.device(coord_target) - zero = xp.zeros((nf, 1, 3), dtype=coord_target.dtype, device=dev) - tokens = xp.concat([zero, atoms, zero], axis=1) + if virtual: + zero = xp.zeros((nf, 1, 3), dtype=coord_target.dtype, device=dev) + tokens = xp.concat([zero, atoms, zero], axis=1) + else: + tokens = atoms if tokens.shape[1] < ncol: pad = xp.zeros( (nf, ncol - tokens.shape[1], 3), dtype=coord_target.dtype, device=dev @@ -180,6 +211,7 @@ def __init__( noise_type: str = "uniform", noise: float = 1.0, data_seed: int = 1, + virtual_tokens: bool = True, **kwargs: float, ) -> None: self.masked_token_loss = masked_token_loss @@ -197,6 +229,52 @@ def __init__( self.noise_type = noise_type self.noise = noise self.data_seed = data_seed + # Uni-Mol wraps each molecule in BOS and EOS and counts them among the + # distance columns. A DPA backbone wraps it in nothing. + self.virtual_tokens = virtual_tokens + + def check_backbone_outputs(self, available: Iterable[str]) -> None: + """Refuse a weighted term that this backbone cannot feed. + + The five weights are independent of which fitting is on the other end, + so a configuration can ask for a term whose model output does not + exist -- the two norm regularisers on a DPA backbone being the case + that arises in practice, since they describe Uni-Mol's own transformer + and a DPA backbone has nothing corresponding. Left alone that surfaces + as a ``KeyError`` on the first batch, after the data has been read and + the statistics computed; this says what to change instead, and the + trainer calls it while the model and the loss are being wired together. + + Parameters + ---------- + available : Iterable[str] + The names the model emits, either from its output definition or + from a batch it has already produced. + + Raises + ------ + ValueError + If any term with a non-zero weight has no matching output. + """ + have = set(available) + missing = [ + (knob, key) + for knob, key in TERM_OUTPUT.items() + if getattr(self, knob) > 0 and key not in have + ] + if not missing: + return + keys = ", ".join(key for _, key in missing) + knobs = ", ".join(knob for knob, _ in missing) + raise ValueError( + f"this backbone emits no {keys}, so the objective cannot evaluate " + f"the term(s) weighted by {knobs}. Uni-Mol's two norm regularisers " + "constrain quantities belonging to its own transformer, and a DPA " + "backbone has no counterpart for them, so on a DPA backbone their " + "weights belong at zero -- as examples/unimol/dpa_pretrain/" + "input.json sets them. Otherwise give the fitting the head that " + f"produces {keys}." + ) def call( self, @@ -208,6 +286,10 @@ def call( ) -> tuple[Array, dict[str, Array]]: """Evaluate the five terms and their weighted sum.""" del learning_rate, natoms, mae + # The trainer already checked this while wiring the two together; it is + # repeated here so that a loss built by hand fails the same way rather + # than on a missing key several lines down. + self.check_backbone_outputs(model_dict.keys()) mask = model_dict.get("mask") token_target = label_dict["unimol_token_target"] xp = array_api_compat.array_namespace(token_target) @@ -244,18 +326,42 @@ def add(term: Array, weight: float, name: str) -> None: "coord_loss", ) if self.masked_dist_loss > 0: - # Rows are the corrupted atoms; columns are every non-padding token, - # BOS, EOS and the diagonal included (losses/unimol.py:159-180). + # Rows are the corrupted atoms; columns are every covered, non-padding + # token. Uni-Mol includes BOS/EOS; a DPA head has only real-atom + # columns and its pair mask selects the neighbour or all-pairs subset. ncol = model_dict["pair_dist"].shape[-1] token_mask = label_dict.get("unimol_token_mask") if token_mask is None: - token_mask = _token_mask_from_atoms(mask, ncol) + token_mask = _token_mask_from_atoms(mask, ncol, self.virtual_tokens) dist_target = label_dict.get("unimol_dist_target") if dist_target is None: dist_target = _clean_distances( - label_dict["unimol_coord_target"], mask, ncol + label_dict["unimol_coord_target"], mask, ncol, self.virtual_tokens + ) + # A head that covers only part of the pair axis says so; Uni-Mol's + # covers all of it and emits nothing here. Which head is on the other + # end also settles whether there are virtual tokens to count, and + # getting that wrong is silent: every label shifts one column, so + # every corrupted row trains against another atom's distances. The + # two are cross-checked here rather than left to the configuration. + covered = model_dict.get("pair_mask") + if covered is None and not self.virtual_tokens: + raise ValueError( + "virtual_tokens=False, but this backbone emits no pair " + "coverage, which means it is the Uni-Mol one -- and that " + "wraps every molecule in BOS and EOS. Scoring it without " + "them shifts every distance label by one column" + ) + if covered is not None and self.virtual_tokens: + raise ValueError( + "virtual_tokens=True, but this backbone wraps no virtual " + "tokens around a molecule (it emits pair coverage, so it is " + "a DPA one). Scoring it as though it did shifts every " + "distance label by one column; set virtual_tokens=False" ) pair_mask = masked[..., None] & xp.astype(token_mask, xp.bool)[:, None, :] + if covered is not None: + pair_mask = pair_mask & xp.astype(covered, xp.bool) dist_label = (dist_target[pair_mask] - DIST_MEAN) / DIST_STD add( _smooth_l1(model_dict["pair_dist"][pair_mask], dist_label, self.beta), @@ -313,7 +419,7 @@ def serialize(self) -> dict: """Serialize the loss module.""" return { "@class": "UniMolLoss", - "@version": 1, + "@version": 2, "masked_token_loss": self.masked_token_loss, "masked_coord_loss": self.masked_coord_loss, "masked_dist_loss": self.masked_dist_loss, @@ -327,12 +433,13 @@ def serialize(self) -> dict: "noise_type": self.noise_type, "noise": self.noise, "data_seed": self.data_seed, + "virtual_tokens": self.virtual_tokens, } @classmethod def deserialize(cls, data: dict) -> "UniMolLoss": """Deserialize the loss module.""" data = data.copy() - check_version_compatibility(data.pop("@version"), 1, 1) + check_version_compatibility(data.pop("@version"), 2, 1) data.pop("@class") return cls(**data) diff --git a/deepmd/dpmodel/model/__init__.py b/deepmd/dpmodel/model/__init__.py index ada9aa58e0..41db932d27 100644 --- a/deepmd/dpmodel/model/__init__.py +++ b/deepmd/dpmodel/model/__init__.py @@ -48,6 +48,9 @@ from .spin_model import ( SpinModel, ) +from .unimol_dpa_pretrain_model import ( + UniMolDPAPretrainModel, +) from .unimol_pretrain_model import ( UniMolPretrainModel, ) @@ -64,6 +67,7 @@ "PolarModel", "PropertyModel", "SpinModel", + "UniMolDPAPretrainModel", "UniMolPretrainModel", "make_model", ] diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py index d49ad1d9df..a281d024fa 100644 --- a/deepmd/dpmodel/model/model_factory.py +++ b/deepmd/dpmodel/model/model_factory.py @@ -54,6 +54,18 @@ def get_model_components( fitting_data["mixed_types"] = descriptor.mixed_types() if fitting_type in {"dipole", "polar"}: fitting_data["embedding_width"] = descriptor.get_dim_emb() + if fitting_type == "unimol_dpa_pretrain": + # The coordinate head projects the backbone's equivariant state, whose + # shape follows the degree that backbone reads out. Reading it here is + # what would otherwise raise first, with an AttributeError naming an + # attribute the reader has no reason to know about. + if not hasattr(descriptor, "node_readout_lmax"): + raise TypeError( + "the unimol DPA objective needs a backbone that reads out an " + "equivariant state, which is what its coordinate head projects; " + f"{type(descriptor).__name__} does not" + ) + fitting_data["node_readout_lmax"] = descriptor.node_readout_lmax fitting_data["dim_descrpt"] = descriptor.get_dim_out() if "direct" in fitting_type: fitting_data["out_dim"] = descriptor.get_dim_emb() diff --git a/deepmd/dpmodel/model/unimol_dpa_pretrain_model.py b/deepmd/dpmodel/model/unimol_dpa_pretrain_model.py new file mode 100644 index 0000000000..3e8a8e60fc --- /dev/null +++ b/deepmd/dpmodel/model/unimol_dpa_pretrain_model.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Model wrapper for Uni-Mol pretraining on a DPA backbone.""" + +from typing import ( + Any, +) + +import array_api_compat + +from deepmd.dpmodel.array_api import ( + Array, +) +from deepmd.dpmodel.atomic_model import ( + DPUniMolDPAAtomicModel, +) +from deepmd.dpmodel.common import ( + NativeOP, +) +from deepmd.dpmodel.model.base_model import ( + BaseModel, +) + +from .dp_model import ( + DPModelCommon, +) +from .make_model import ( + make_model, +) + + +def _reject_periodic(box) -> None: # noqa: ANN001 + """Refuse a cell, with an explanation rather than a quiet mistraining.""" + if box is None: + return + xp = array_api_compat.array_namespace(box) + if bool(xp.any(box != 0)): + raise ValueError( + "the unimol objective does not support periodic boundaries: its " + "distance target is built from plain coordinate differences with no " + "minimum-image convention, and its coverage keeps only local " + "neighbours, so a cell would train against wrong labels rather than " + "fail. Pass box=None" + ) + + +DPUniMolDPAPretrainModel_ = make_model( + DPUniMolDPAAtomicModel, T_Bases=(NativeOP, BaseModel) +) + + +@BaseModel.register("unimol_dpa_pretrain") +class UniMolDPAPretrainModel(DPModelCommon, DPUniMolDPAPretrainModel_): + r"""Uni-Mol's pretraining objective on a DPA backbone. + + Predicts the element of every corrupted atom, denoises the coordinates from + the backbone's equivariant state, and predicts the clean pairwise distances + from pairs of atom representations. Nothing reduces to a frame total and + nothing is differentiated with respect to the coordinates: this is + representation learning, not a potential energy surface. + + Periodic frames are refused. The backbone handles them, but this objective + does not: the distance target is built from plain coordinate differences + with no minimum-image convention, and the coverage mask keeps only local + neighbours, so under a cell the labels are wrong by up to several Angstrom + and some exceed the cut-off entirely. It would run and quietly mistrain. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + DPModelCommon.__init__(self) + DPUniMolDPAPretrainModel_.__init__(self, *args, **kwargs) + + def call( + self, + coord: Array, + atype: Array, + box: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, + do_atomic_virial: bool = False, + charge_spin: Array | None = None, + ) -> dict[str, Array]: + """Evaluate the heads on a frame. + + Raises + ------ + ValueError + If a periodic cell is supplied; see the class docstring. + + Parameters + ---------- + coord : Array + Atomic coordinates, flattened per frame or shaped as ``(nf, nall, 3)``. + atype : Array + Atomic type indices. + box : Array, optional + Periodic cell, which must be absent or all zero for this objective. + fparam, aparam : Array, optional + Unused frame and atomic parameters. + do_atomic_virial : bool + Whether to request atomic virials; this objective does not expose + coordinate derivatives. + charge_spin : Array, optional + Optional charge/spin conditioning passed to the DPA backbone. + + Returns + ------- + dict[str, Array] + Element logits, clean-coordinate predictions, pair distances and + pair coverage, with absent optional heads omitted. + """ + _reject_periodic(box) + model_ret = self.call_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, + ) + return {k: v for k, v in model_ret.items() if v is not None} diff --git a/deepmd/pt_expt/fitting/__init__.py b/deepmd/pt_expt/fitting/__init__.py index 63bcc6c2ce..bd875950c5 100644 --- a/deepmd/pt_expt/fitting/__init__.py +++ b/deepmd/pt_expt/fitting/__init__.py @@ -23,6 +23,9 @@ from .property_fitting import ( PropertyFittingNet, ) +from .unimol_dpa_pretrain import ( + UniMolDPAPretrainFitting, +) from .unimol_pretrain import ( UniMolPretrainFitting, ) @@ -36,5 +39,6 @@ "PolarFitting", "PropertyFittingNet", "SeZMEnergyFittingNet", + "UniMolDPAPretrainFitting", "UniMolPretrainFitting", ] diff --git a/deepmd/pt_expt/fitting/unimol_dpa_pretrain.py b/deepmd/pt_expt/fitting/unimol_dpa_pretrain.py new file mode 100644 index 0000000000..941ea09a3e --- /dev/null +++ b/deepmd/pt_expt/fitting/unimol_dpa_pretrain.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +from typing import ( + Any, +) + +import torch + +from deepmd.dpmodel.fitting.unimol_dpa_pretrain import ( + UniMolDPAPretrainFitting as UniMolDPAPretrainFittingDP, +) +from deepmd.pt_expt.common import ( + torch_module, +) + +from .base_fitting import ( + BaseFitting, +) + +# The coordinate head projects with the descriptor's own ``SO3Linear``, which +# keeps its weights as plain arrays. On this backend ``dpmodel_setattr`` +# registers a plain array as a *buffer*, so without this the head would hold no +# parameters at all: the optimizer is built from ``wrapper.parameters()``, so it +# would train as a frozen random projection, and silently -- the coordinate term +# still falls, because gradients reach the backbone through it either way. The +# DPA4 descriptors each carry their own copy of this promotion for the same +# reason; this is the fitting's. +_TRAINABLE_ATTRS: dict[str, tuple[str, ...]] = { + "SO3Linear": ("weight", "bias"), +} + + +def _promote_trainable_tree(module: torch.nn.Module) -> torch.nn.Module: + """Re-register the float buffers that are meant to be learned. + + Runs after the tree is built: deserialization assigns arrays onto nested + attributes, which would be registered as buffers again. + """ + for sub in module.modules(): + if not getattr(sub, "trainable", True): + continue + for name in _TRAINABLE_ATTRS.get(type(sub).__name__, ()): + buf = sub._buffers.get(name) + if buf is None or not buf.is_floating_point(): + continue + del sub._buffers[name] + setattr(sub, name, torch.nn.Parameter(buf, requires_grad=True)) + return module + + +@BaseFitting.register("unimol_dpa_pretrain") +@torch_module +class UniMolDPAPretrainFitting(UniMolDPAPretrainFittingDP): + """Uni-Mol's heads on a DPA backbone, PyTorch-Exportable backend.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + _promote_trainable_tree(self) + + @classmethod + def deserialize(cls, data: dict) -> "UniMolDPAPretrainFitting": + """Deserialize the fitting, keeping its learned weights learnable.""" + return _promote_trainable_tree(super().deserialize(data)) diff --git a/deepmd/pt_expt/model/__init__.py b/deepmd/pt_expt/model/__init__.py index 22a1c2e7db..78bd29abbe 100644 --- a/deepmd/pt_expt/model/__init__.py +++ b/deepmd/pt_expt/model/__init__.py @@ -42,6 +42,9 @@ from .spin_ener_model import ( SpinEnergyModel, ) +from .unimol_dpa_pretrain_model import ( + UniMolDPAPretrainModel, +) from .unimol_pretrain_model import ( UniMolPretrainModel, ) @@ -59,6 +62,7 @@ "PolarModel", "PropertyModel", "SpinEnergyModel", + "UniMolDPAPretrainModel", "UniMolPretrainModel", "get_model", "make_hessian_model", diff --git a/deepmd/pt_expt/model/unimol_dpa_pretrain_model.py b/deepmd/pt_expt/model/unimol_dpa_pretrain_model.py new file mode 100644 index 0000000000..5ca07ecd58 --- /dev/null +++ b/deepmd/pt_expt/model/unimol_dpa_pretrain_model.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import torch + +from deepmd.dpmodel.atomic_model import ( + DPUniMolDPAAtomicModel, +) +from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, +) + +from .make_model import ( + make_model, +) +from .model import ( + BaseModel, +) + +DPUniMolDPAPretrainModel_ = make_model(DPUniMolDPAAtomicModel, T_Bases=(BaseModel,)) + + +@BaseModel.register("unimol_dpa_pretrain") +class UniMolDPAPretrainModel(DPModelCommon, DPUniMolDPAPretrainModel_): + """Uni-Mol's objective on a DPA backbone, PyTorch-Exportable backend.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + DPModelCommon.__init__(self) + DPUniMolDPAPretrainModel_.__init__(self, *args, **kwargs) + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Evaluate the heads on a frame; a periodic cell is refused.""" + if box is not None and bool(torch.any(box != 0)): + raise ValueError( + "the unimol objective does not support periodic boundaries: its " + "distance target is built from plain coordinate differences with " + "no minimum-image convention, and its coverage keeps only local " + "neighbours, so a cell would train against wrong labels rather " + "than fail. Pass box=None" + ) + model_ret = self.forward_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, + ) + return {k: v for k, v in model_ret.items() if v is not None} diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 9130f5753b..199d4025b6 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -363,8 +363,14 @@ def get_loss( return PropertyLoss(**loss_params) elif loss_type == "unimol": # Self-supervised: it takes no learning rate and no model geometry, - # because its targets come from the corruption it defines itself. - return UniMolLoss(**loss_params) + # because its targets come from the corruption it defines itself. It + # does need to know which heads the model carries: the five term + # weights are set independently of the fitting, so a configuration can + # weight a term the backbone has no output for. Checked here, while the + # two are being wired together, rather than on the first batch. + loss = UniMolLoss(**loss_params) + loss.check_backbone_outputs(_model.model_output_def().keys()) + return loss else: raise ValueError(f"Unsupported loss type for pt_expt: {loss_type}") diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 982c5d2a26..c3bd0ee53b 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -2781,6 +2781,14 @@ def descrpt_se_a_mask_args() -> list[Argument]: "The three self-supervised heads of Uni-Mol v1 pretraining: element prediction, coordinate " "denoising and pairwise distance prediction." ) +doc_fitting_unimol_dpa_pretrain = ( + "Uni-Mol v1's three self-supervised heads reading a DPA backbone. The element head takes the " + "per-atom representation as in Uni-Mol; the coordinate head takes the backbone's equivariant " + "state, because a DPA backbone carries no pair channel to read an update from; and the " + "distance head describes a pair by its two endpoints for the same reason. Uni-Mol's two norm " + "regularisers have no counterpart here, since they constrain quantities belonging to its own " + "transformer." +) @descrpt_args_plugin.register( @@ -3530,6 +3538,71 @@ def fitting_unimol_pretrain() -> list[Argument]: ] +@fitting_args_plugin.register( + "unimol_dpa_pretrain", + doc=supported_backends("pt_expt") + doc_fitting_unimol_dpa_pretrain, +) +def fitting_unimol_dpa_pretrain() -> list[Argument]: + doc_seed = "Random seed for parameter initialization" + doc_precision = f"The precision of the parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." + doc_n_token = ( + "Size of the Uni-Mol vocabulary, which is the width of the element head." + ) + doc_max_atoms = ( + "Largest molecule accepted, which fixes the distance head's column count. A frame with " + "more atoms than this is refused rather than truncated." + ) + doc_dist_hidden = "Width of the distance head's hidden layer." + doc_dist_coverage = ( + "Which atom pairs the distance term covers. 'neighbour' follows the backbone's own " + "neighbour list for the scored entries, but the dense pair head still evaluates all " + "nloc^2 pairs, so its head compute remains O(nloc^2); it reuses the locality the rest of " + "deepmd trains on, but it sees only the pairs inside the cut-off -- on drug-like molecules " + "a 6 A cut-off holds roughly half of all pairs, and less for the larger ones. 'all_pairs' covers " + "every pair, which is Uni-Mol's own coverage and what to use to reproduce its training; it " + "costs O(nloc^2) and does not share the backbone's neighbour structure. The setting " + "changes which pairs are scored and nothing else: the head predicts the same numbers " + "either way." + ) + doc_activation_function = f"The activation function of the heads. Supported: {list_to_doc(ACTIVATION_FN_DICT.keys())}" + doc_mask_token_head = "Build the element-prediction head." + doc_coord_head = ( + "Build the coordinate-denoising head. It reads the backbone's equivariant state, so the " + "backbone must read out at least degree 1." + ) + doc_dist_head = "Build the distance-prediction head." + return [ + Argument("n_token", int, optional=True, default=31, doc=doc_n_token), + Argument("max_atoms", int, optional=True, default=256, doc=doc_max_atoms), + Argument("dist_hidden", int, optional=True, default=64, doc=doc_dist_hidden), + Argument( + "dist_coverage", + str, + optional=True, + default="neighbour", + doc=doc_dist_coverage, + ), + Argument( + "activation_function", + str, + optional=True, + default="gelu_erf", + doc=doc_activation_function, + ), + Argument( + "mask_token_head", + bool, + optional=True, + default=True, + doc=doc_mask_token_head, + ), + Argument("coord_head", bool, optional=True, default=True, doc=doc_coord_head), + Argument("dist_head", bool, optional=True, default=True, doc=doc_dist_head), + Argument("precision", str, optional=True, default="default", doc=doc_precision), + Argument("seed", [int, list, None], optional=True, doc=doc_seed), + ] + + def fitting_variant_type_args() -> Variant: doc_descrpt_type = "The type of the fitting." @@ -5605,8 +5678,16 @@ def loss_unimol() -> list[Argument]: doc_masked_token_loss = "Weight of the element-prediction term." doc_masked_coord_loss = "Weight of the coordinate-denoising term." doc_masked_dist_loss = "Weight of the distance-prediction term." - doc_x_norm_loss = "Weight of the node-norm regularizer." - doc_delta_pair_repr_norm_loss = "Weight of the pair-delta-norm regularizer." + doc_x_norm_loss = ( + "Weight of the node-norm regularizer. It constrains a quantity of " + "Uni-Mol's own transformer, so only the unimol_pretrain fitting " + "produces it; on any other backbone this must be 0." + ) + doc_delta_pair_repr_norm_loss = ( + "Weight of the pair-delta-norm regularizer. As with the node-norm " + "regularizer, only the unimol_pretrain fitting produces it; on any " + "other backbone this must be 0." + ) doc_beta = ( "Transition point of the smooth L1 used by the coordinate and distance terms." ) @@ -5626,6 +5707,11 @@ def loss_unimol() -> list[Argument]: "data (DP_LMDB_NUM_WORKERS=0); with decoder workers the draws follow how " "frames were distributed." ) + doc_virtual_tokens = ( + "Whether the backbone wraps each molecule in two virtual tokens and counts them among " + "the distance columns, which Uni-Mol's own backbone does and a DPA backbone does not. " + "Scoring a backbone with the wrong setting compares against columns that are not there." + ) return [ Argument( "masked_token_loss", @@ -5689,6 +5775,13 @@ def loss_unimol() -> list[Argument]: ), Argument("noise", [float, int], optional=True, default=1.0, doc=doc_noise), Argument("data_seed", int, optional=True, default=1, doc=doc_data_seed), + Argument( + "virtual_tokens", + bool, + optional=True, + default=True, + doc=doc_virtual_tokens, + ), ] diff --git a/deepmd/utils/unimol_data.py b/deepmd/utils/unimol_data.py index 712997e7d0..e54151b29e 100644 --- a/deepmd/utils/unimol_data.py +++ b/deepmd/utils/unimol_data.py @@ -95,7 +95,14 @@ def _conformers(record: dict[str, Any], add_2d_conformer: bool) -> list[np.ndarr AllChem, ) - mol = Chem.AddHs(Chem.MolFromSmiles(record["smi"])) + parsed = Chem.MolFromSmiles(record["smi"]) + if parsed is None: + # RDKit returns None for a SMILES it cannot parse, and AddHs would + # raise on it, aborting the whole conversion. The 3D conformers + # are still good, so the molecule keeps them and goes without the + # 2D one. + return conformers + mol = Chem.AddHs(parsed) AllChem.Compute2DCoords(mol) coords = mol.GetConformer().GetPositions().astype(np.float32) conformers.append(coords[: len(record["atoms"])]) diff --git a/doc/model/index.rst b/doc/model/index.rst index c6db98d169..711e26ad9c 100644 --- a/doc/model/index.rst +++ b/doc/model/index.rst @@ -14,6 +14,7 @@ Model dpa4 dpa4c unimol + unimol-dpa train-hybrid sel train-energy diff --git a/doc/model/unimol-dpa.md b/doc/model/unimol-dpa.md new file mode 100644 index 0000000000..c73d5c4e9a --- /dev/null +++ b/doc/model/unimol-dpa.md @@ -0,0 +1,128 @@ +# Uni-Mol pretraining on a DPA backbone {{ pytorch_icon }} {{ dpmodel_icon }} + +> [!NOTE] +> **Supported backends**: PyTorch-Exportable {{ pytorch_icon }}, DP {{ dpmodel_icon }} + +Uni-Mol's self-supervised objective does not need Uni-Mol's transformer. This +runs the same three tasks -- recover the element of a corrupted atom, denoise +the coordinates, predict the clean pairwise distances -- on a DPA4 backbone, so +that molecular pretraining data can train the same descriptor that DFT data +trains. + +See [the Uni-Mol backbone](unimol.md) for the objective itself: the corruption, +the five terms and their weights, and how the data is prepared. Only what the +heads read differs here. + +## What the heads read + +Uni-Mol's transformer carries a representation per atom pair and wraps each +molecule in two virtual tokens. DPA4 carries neither: it attends per edge with a +scatter softmax and returns no pair axis at all. So: + +| Head | Uni-Mol reads | Here it reads | +| ---------- | --------------------------- | ------------------------------------- | +| element | the per-atom representation | the same | +| coordinate | the pair channel | the l=1 part of the equivariant state | +| distance | the pair channel | the two endpoints' representations | + +The coordinate head projects the backbone's equivariant state with a degree-wise +linear whose weights are shared across the three $m$ components, so its +predicted displacement rotates with the molecule rather than merely being three +numbers. The displacement is added to the input coordinates before the +coordinate loss, matching Uni-Mol's `coord + update` head and making the clean +coordinate prediction translate with the input. The distance head combines the +two endpoints symmetrically, $[h_i + h_j \,\|\, h_i \odot h_j]$, which makes the +predicted matrix symmetric by construction. + +Uni-Mol's two norm regularisers constrain quantities belonging to its own +transformer and have no counterpart here, so they carry no weight; and because +there are no virtual tokens, the objective is configured with +`virtual_tokens: false`. + +Both are weighted `0.01` by default, which is Uni-Mol's own recipe, so a +configuration that selects this fitting and leaves the loss alone asks for two +terms this backbone cannot produce. That is refused when the model and the loss +are wired together, naming the weights to zero -- so the two lines below are +required, not decorative: + +```json +"loss": { + "type": "unimol", + "x_norm_loss": 0.0, + "delta_pair_repr_norm_loss": 0.0, + "virtual_tokens": false +} +``` + +## Which pairs the distance term covers + +This is the one place the objective departs from Uni-Mol by construction, and it +is a choice rather than a limitation of the port. + +`dist_coverage: neighbour` (default) +: Only pairs inside the backbone's neighbour list. The mask reuses the locality +the rest of deepmd trains on and reduces the scored entries, but the dense +pair head still evaluates all `nloc²` pairs, so its head compute remains +`O(nloc²)`. It also sees less than Uni-Mol does: on drug-like molecules a 6 Å +cut-off holds about half of all pairs, and under a third for the largest. + +`dist_coverage: all_pairs` +: Every pair, which is Uni-Mol's own coverage and what to use to reproduce its +training. It costs `O(nloc^2)` and does not share the backbone's neighbour +structure. + +The setting selects which pairs are scored and nothing else: the head predicts +the same numbers either way. `all_pairs` includes the diagonal, because upstream +scores the zero self-distance too; a neighbour list never lists an atom as its +own neighbour, so given a list that already holds every other atom the two +differ by exactly those self-pairs and nothing more. + +## The `[MASK]` pseudo-element + +A corrupted atom is carried as `[MASK]`, so the `type_map` declares it alongside +the elements and the backbone's type embedding gains a row for it. Two +consequences are worth knowing before pairing this with an existing model: + +- A DPA4 model pretrained without `[MASK]` has a different `type_map`, and DPA4 + does not implement `change_type_map`, so it cannot be adapted to this one. +- Electronic-configuration type embedding would reject `[MASK]`, since it is + not a real element. DPA4 has no such option, so the two cannot meet today; + worth knowing if this objective is ever put on a backbone that does. + +## Periodic frames + +Refused. The backbone handles a cell; this objective does not. Its distance +target is a plain coordinate difference with no minimum-image convention, and +its coverage keeps only local neighbours, so a periodic frame would train +against labels that are wrong by several Angstrom -- some of them longer than +the cut-off -- rather than fail. The refusal is at the atomic model, so it holds +on the evaluation path too, not only when a box is passed to the model. + +## Reproducibility of the coordinate head + +On the torch backend the equivariant state the coordinate head reads is not +bit-reproducible between identical calls when the backbone runs in single +precision with `use_env_seed` on, which is the shipped default. The difference +measures around 7e-09 — one part in a hundred million, far below anything the +objective resolves — and the element and distance heads are unaffected, since +the scalar read-out is exact. It disappears in double precision or with +`use_env_seed` off. This is a property of the backbone rather than of the heads; +it is mentioned because comparing `coord_update` across two runs will otherwise +look like a bug. + +## Multi-task training + +Not yet. Training one shared backbone from both a DFT branch and this objective +is the point of putting the objective on a DPA backbone, and the multi-task +machinery itself is ready for it -- but DPA4 does not implement `share_params` +on these backends, so the branches cannot be linked. Until it does, this +objective trains a backbone of its own. + +## Training + +```sh +dp --pt-expt train examples/unimol/dpa_pretrain/input.json +``` + +The dataset is the same LMDB one the Uni-Mol backbone uses, built the same way; +see [the Uni-Mol page](unimol.md). diff --git a/examples/unimol/dpa_pretrain/input.json b/examples/unimol/dpa_pretrain/input.json new file mode 100644 index 0000000000..e15bce16e1 --- /dev/null +++ b/examples/unimol/dpa_pretrain/input.json @@ -0,0 +1,116 @@ +{ + "_comment": "Uni-Mol v1's self-supervised objective on a DPA4 backbone. The element head reads the per-atom representation as Uni-Mol's does; the coordinate head reads the backbone's equivariant state, because DPA4 carries no pair channel; the distance head describes a pair by its two endpoints, for the same reason.", + "model": { + "type_map": [ + "C", + "N", + "O", + "S", + "H", + "Cl", + "F", + "Br", + "I", + "Si", + "P", + "B", + "Na", + "K", + "Al", + "Ca", + "Sn", + "As", + "Hg", + "Fe", + "Zn", + "Cr", + "Se", + "Gd", + "Au", + "Li", + "[MASK]" + ], + "_comment_type_map": "[MASK] is the pseudo-element a corrupted atom is carried as; the objective needs it", + "descriptor": { + "type": "dpa4", + "rcut": 6.0, + "channels": 32, + "n_radial": 16, + "edge_norm": false, + "use_env_seed": true, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "mixing_layers": 3, + "radial_so2_mode": "degree_channel", + "radial_so2_rank": 1, + "n_focus": 1, + "focus_dim": 0, + "n_atten_head": 1, + "message_node_so3": true, + "ffn_neurons": 0, + "ffn_so3_grid": true, + "grid_mlp": false, + "grid_branch": [ + 0, + 0, + 1 + ], + "ffn_blocks": 1, + "so3_readout": "mlp", + "use_amp": true, + "precision": "float32", + "seed": 42 + }, + "fitting_net": { + "type": "unimol_dpa_pretrain", + "n_token": 31, + "max_atoms": 256, + "dist_hidden": 64, + "dist_coverage": "neighbour", + "_comment_dist_coverage": "follow the backbone's neighbour list, which is what the rest of deepmd trains on; set all_pairs to cover every pair as Uni-Mol does, at quadratic cost and without sharing that locality", + "seed": 1 + } + }, + "learning_rate": { + "type": "wsd", + "start_lr": 0.0001, + "stop_lr": 1e-12, + "warmup_steps": 10000, + "decay_type": "linear", + "decay_phase_ratio": 1.0 + }, + "loss": { + "type": "unimol", + "masked_token_loss": 1.0, + "masked_coord_loss": 5.0, + "masked_dist_loss": 10.0, + "x_norm_loss": 0.0, + "delta_pair_repr_norm_loss": 0.0, + "virtual_tokens": false, + "_comment": "Uni-Mol's two norm regularisers constrain quantities belonging to its own transformer and have no counterpart here, so they carry no weight; a DPA backbone also wraps no virtual tokens around a molecule" + }, + "training": { + "training_data": { + "systems": "./unimol_train", + "batch_size": 16, + "_comment": "an LMDB dataset path, given as a string; build it with deepmd.utils.unimol_data" + }, + "validation_data": { + "systems": "./unimol_valid", + "batch_size": 16 + }, + "numb_steps": 1000000, + "seed": 1, + "disp_file": "lcurve.out", + "disp_freq": 100, + "save_freq": 10000 + }, + "optimizer": { + "type": "Adam", + "adam_beta1": 0.9, + "adam_beta2": 0.99, + "adam_eps": 1e-06, + "weight_decay": 0.0001 + } +} diff --git a/source/tests/common/dpmodel/test_unimol.py b/source/tests/common/dpmodel/test_unimol.py index 65830dc747..a66e27bf05 100644 --- a/source/tests/common/dpmodel/test_unimol.py +++ b/source/tests/common/dpmodel/test_unimol.py @@ -602,10 +602,15 @@ def test_derived_labels_match_explicit_ones(self) -> None: np.testing.assert_allclose(float(derived), float(explicit), rtol=1e-7) def test_serialize_round_trip(self) -> None: - loss = UniMolLoss(masked_coord_loss=3.0, beta=0.5) + loss = UniMolLoss(masked_coord_loss=3.0, beta=0.5, virtual_tokens=False) clone = UniMolLoss.deserialize(loss.serialize()) self.assertEqual(clone.masked_coord_loss, 3.0) self.assertEqual(clone.beta, 0.5) + self.assertFalse(clone.virtual_tokens) + legacy = loss.serialize() + legacy["@version"] = 1 + legacy.pop("virtual_tokens") + self.assertTrue(UniMolLoss.deserialize(legacy).virtual_tokens) class TestUniMolExample(unittest.TestCase): diff --git a/source/tests/common/dpmodel/test_unimol_data.py b/source/tests/common/dpmodel/test_unimol_data.py index 6a885e7cb1..439f1523a9 100644 --- a/source/tests/common/dpmodel/test_unimol_data.py +++ b/source/tests/common/dpmodel/test_unimol_data.py @@ -4,7 +4,9 @@ import os import pickle import shutil +import sys import tempfile +import types import unittest import unittest.mock @@ -23,6 +25,7 @@ make_unimol_data_transform, ) from deepmd.utils.unimol_data import ( + _conformers, convert_unimol_lmdb, read_unimol_lmdb, ) @@ -377,6 +380,33 @@ def test_limits_are_respected(self) -> None: self.assertEqual(counts["molecules"], 1) self.assertEqual(counts["frames"], 2) + def test_an_unparsable_smiles_only_loses_its_2d_conformer(self) -> None: + """A failed SMILES parse keeps the existing conformers. + + Mocking the optional RDKit dependency keeps this regression test active + in the normal test environment, where RDKit is not installed. Passing + ``None`` to ``AddHs`` used to abort conversion; the 3D conformers do not + depend on the SMILES and must remain available. + """ + rng = np.random.default_rng(1) + conformers = [rng.normal(size=(4, 3)).astype(np.float32) for _ in range(2)] + record = {"atoms": ["C"] * 4, "coordinates": conformers, "smi": "C1CCC"} + chem = types.ModuleType("rdkit.Chem") + chem.MolFromSmiles = unittest.mock.Mock(return_value=None) + chem.AddHs = unittest.mock.Mock(side_effect=AssertionError("must not run")) + chem.AllChem = types.SimpleNamespace(Compute2DCoords=unittest.mock.Mock()) + rdkit = types.ModuleType("rdkit") + rdkit.Chem = chem + with unittest.mock.patch.dict( + sys.modules, {"rdkit": rdkit, "rdkit.Chem": chem} + ): + kept = _conformers(record, add_2d_conformer=True) + self.assertEqual(len(kept), 2) + for actual, expected in zip(kept, conformers, strict=True): + np.testing.assert_array_equal(actual, expected) + chem.MolFromSmiles.assert_called_once_with("C1CCC") + chem.AddHs.assert_not_called() + if __name__ == "__main__": unittest.main() diff --git a/source/tests/common/dpmodel/test_unimol_dpa.py b/source/tests/common/dpmodel/test_unimol_dpa.py new file mode 100644 index 0000000000..a087fe77ac --- /dev/null +++ b/source/tests/common/dpmodel/test_unimol_dpa.py @@ -0,0 +1,595 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Uni-Mol pretraining heads on a DPA backbone.""" + +import unittest + +import numpy as np + +from deepmd.dpmodel.atomic_model import ( + DPUniMolDPAAtomicModel, +) +from deepmd.dpmodel.descriptor.dpa4 import ( + DescrptDPA4, +) +from deepmd.dpmodel.fitting.unimol_dpa_heads import ( + EquivariantCoordHead, + PairDistanceHead, + l1_to_cartesian, +) +from deepmd.dpmodel.fitting.unimol_dpa_pretrain import ( + UniMolDPAPretrainFitting, +) +from deepmd.dpmodel.loss.unimol import ( + UniMolLoss, +) + +SMALL = {"nloc": 6, "nnei": 5, "ntypes": 2} + + +def _rotation(rng): + """A random proper rotation.""" + q, _ = np.linalg.qr(rng.normal(size=(3, 3))) + if np.linalg.det(q) < 0: + q[:, 0] *= -1 + return q + + +class TestEquivariantCoordHead(unittest.TestCase): + """The head that replaces Uni-Mol's pair-channel coordinate update.""" + + def setUp(self) -> None: + self.rng = np.random.default_rng(0) + nloc, nnei = SMALL["nloc"], SMALL["nnei"] + # float64 throughout: the property under test is exact, and the default + # single precision would only show it to about 1e-6. + self.desc = DescrptDPA4( + rcut=6.0, + rcut_smth=5.0, + sel=nnei, + ntypes=SMALL["ntypes"], + seed=0, + precision="float64", + ) + self.head = EquivariantCoordHead( + lmax=self.desc.node_readout_lmax, + channels=self.desc.channels, + precision="float64", + seed=1, + ) + self.coord = self.rng.normal(size=(1, nloc, 3)) * 1.5 + self.atype = self.rng.integers(0, SMALL["ntypes"], size=(1, nloc)) + self.nlist = np.stack( + [ + np.stack( + [ + np.array([j for j in range(nloc) if j != i][:nnei]) + for i in range(nloc) + ] + ) + ] + ) + self.mapping = np.tile(np.arange(nloc), (1, 1)) + + def _predict(self, coord): + _, latent = self.desc.call_with_latent( + coord, self.atype, self.nlist, mapping=self.mapping + ) + return self.head(latent) + + def test_one_vector_per_atom(self) -> None: + out = self._predict(self.coord) + self.assertEqual(out.shape, (SMALL["nloc"], 3)) + self.assertTrue(bool(np.all(np.isfinite(out)))) + + def test_rotating_the_molecule_rotates_the_prediction(self) -> None: + """The point of reading l=1 rather than three arbitrary numbers. + + A coordinate update has to be equivariant; if it were not, the model + could not denoise a rotated copy of a structure it had already learned. + """ + base = self._predict(self.coord) + for trial in range(5): + rot = _rotation(self.rng) + with self.subTest(rotation=trial): + turned = self._predict(self.coord @ rot.T) + np.testing.assert_allclose(turned, base @ rot.T, rtol=1e-9, atol=1e-9) + + def test_translating_the_molecule_changes_nothing(self) -> None: + base = self._predict(self.coord) + shifted = self._predict(self.coord + self.rng.normal(size=3)) + np.testing.assert_allclose(shifted, base, rtol=1e-9, atol=1e-9) + + def test_reading_the_rows_as_xyz_would_not_be_equivariant(self) -> None: + """Guards the decode itself, not just the head around it. + + SeZM packs l=1 in its own basis, so taking the rows as (x, y, z) gives + something that does not rotate. This pins the mapping that does. + """ + rows = self.rng.normal(size=(4, 3)) + np.testing.assert_allclose( + l1_to_cartesian(rows), + np.stack([-rows[:, 2], rows[:, 0], rows[:, 1]], axis=-1), + ) + + def test_a_backbone_without_l1_is_refused(self) -> None: + with self.assertRaisesRegex(ValueError, "at least degree 1"): + EquivariantCoordHead(lmax=0, channels=8) + + def test_serialize_round_trip(self) -> None: + revived = EquivariantCoordHead.deserialize(self.head.serialize()) + _, latent = self.desc.call_with_latent( + self.coord, self.atype, self.nlist, mapping=self.mapping + ) + np.testing.assert_allclose(revived(latent), self.head(latent)) + + +class TestPairDistanceHead(unittest.TestCase): + """The head that replaces Uni-Mol's pair-channel distance prediction.""" + + def setUp(self) -> None: + rng = np.random.default_rng(0) + self.nf, self.nloc, self.nnei, self.dim = 2, 5, 3, 8 + self.node = rng.normal(size=(self.nf, self.nloc, self.dim)) + self.nlist = np.full((self.nf, self.nloc, self.nnei), -1, dtype=np.int64) + for f in range(self.nf): + for i in range(self.nloc): + others = [j for j in range(self.nloc) if j != i][:2] + self.nlist[f, i, : len(others)] = others + + def _head(self, coverage): + return PairDistanceHead( + dim_descrpt=self.dim, hidden=6, coverage=coverage, seed=0 + ) + + def test_a_distance_is_symmetric(self) -> None: + """Not approximately: the pair feature is symmetric in the endpoints.""" + for coverage in PairDistanceHead.COVERAGES: + with self.subTest(coverage=coverage): + dist, _ = self._head(coverage)(self.node, self.nlist) + np.testing.assert_allclose( + dist, + np.transpose(dist, (0, 2, 1)), + rtol=0.0, + atol=1e-14, + ) + + def test_the_diagonal_follows_the_coverage(self) -> None: + """Upstream scores the zero self-distance; a neighbour list cannot. + + An atom never lists itself as a neighbour, so ``neighbour`` excludes the + diagonal by construction. ``all_pairs`` claims to reproduce Uni-Mol's + coverage, and Uni-Mol takes every non-padding column, diagonal included + -- those entries carry the most extreme normalised label in the set. + """ + _, near = self._head("neighbour")(self.node, self.nlist) + np.testing.assert_array_equal( + np.einsum("fii->fi", near), np.zeros((self.nf, self.nloc)) + ) + _, whole = self._head("all_pairs")(self.node, self.nlist) + np.testing.assert_array_equal( + np.einsum("fii->fi", whole), np.ones((self.nf, self.nloc)) + ) + + def test_coverage_selects_pairs_and_nothing_else(self) -> None: + """The option changes which pairs the objective sees, not the model. + + Uni-Mol's distance term covers every pair; a neighbour list covers a + subset. Both heads predict the same numbers from the same weights, so + the choice is purely one of coverage. + """ + near, whole = self._head("neighbour"), self._head("all_pairs") + d_near, m_near = near(self.node, self.nlist) + d_whole, m_whole = whole(self.node, self.nlist) + + np.testing.assert_allclose(d_near, d_whole) + self.assertTrue(bool(np.all(m_near <= m_whole))) + # all_pairs is every pair including the diagonal, as upstream scores it + self.assertEqual(int(m_whole.sum()), self.nf * self.nloc * self.nloc) + self.assertLess(int(m_near.sum()), int(m_whole.sum())) + # and it is exactly the neighbour list, padding excluded + expected = np.zeros_like(m_near) + for f in range(self.nf): + for i in range(self.nloc): + for j in self.nlist[f, i]: + if j >= 0: + expected[f, i, j] = 1.0 + np.testing.assert_array_equal(m_near, expected) + + def test_padding_in_the_neighbour_list_covers_nothing(self) -> None: + """A padded entry is negative and would wrap if used as an index.""" + empty = np.full_like(self.nlist, -1) + _, mask = self._head("neighbour")(self.node, empty) + np.testing.assert_array_equal(mask, np.zeros_like(mask)) + + def test_an_unknown_coverage_is_refused(self) -> None: + with self.assertRaisesRegex(ValueError, "unknown coverage"): + PairDistanceHead(dim_descrpt=self.dim, coverage="everything") + + def test_serialize_round_trip(self) -> None: + head = self._head("neighbour") + revived = PairDistanceHead.deserialize(head.serialize()) + a, ma = head(self.node, self.nlist) + b, mb = revived(self.node, self.nlist) + np.testing.assert_allclose(a, b) + np.testing.assert_array_equal(ma, mb) + + +class TestUniMolDPAPretrainFitting(unittest.TestCase): + """The three heads assembled on a real DPA4 backbone.""" + + def setUp(self) -> None: + rng = np.random.default_rng(0) + self.nloc, nnei, ntypes = 6, 5, 3 + self.max_atoms = 10 + self.desc = DescrptDPA4( + rcut=6.0, + rcut_smth=5.0, + sel=nnei, + ntypes=ntypes, + seed=0, + precision="float64", + ) + self.fitting = UniMolDPAPretrainFitting( + ntypes=ntypes, + dim_descrpt=self.desc.channels, + node_readout_lmax=self.desc.node_readout_lmax, + max_atoms=self.max_atoms, + dist_hidden=16, + precision="float64", + seed=1, + ) + self.coord = rng.normal(size=(1, self.nloc, 3)) * 1.5 + self.atype = rng.integers(0, ntypes, size=(1, self.nloc)) + self.nlist = np.stack( + [ + np.stack( + [ + np.array([j for j in range(self.nloc) if j != i][:nnei]) + for i in range(self.nloc) + ] + ) + ] + ) + + def _run(self): + node, latent = self.desc.call_with_latent( + self.coord, + self.atype, + self.nlist, + mapping=np.tile(np.arange(self.nloc), (1, 1)), + ) + return self.fitting.call_atoms(node, latent, self.nlist, self.coord) + + def test_every_output_is_declared(self) -> None: + """The output machinery indexes the definition for every key returned. + + A key the definition does not declare raises KeyError on the way out of + the model, so returning one is not a harmless extra. + """ + out = self._run() + self.assertEqual(set(out), set(self.fitting.output_def().var_defs)) + + def test_shapes(self) -> None: + out = self._run() + self.assertEqual(out["token_logits"].shape, (1, self.nloc, 31)) + self.assertEqual(out["coord_update"].shape, (1, self.nloc, 3)) + self.assertEqual(out["pair_dist"].shape, (1, self.nloc, self.max_atoms)) + self.assertEqual(out["pair_mask"].shape, (1, self.nloc, self.max_atoms)) + for name, value in out.items(): + with self.subTest(output=name): + self.assertTrue(bool(np.all(np.isfinite(value)))) + + def test_coordinate_prediction_translates_with_the_input(self) -> None: + """The head predicts ``coord + update`` like Uni-Mol v1.""" + base = self._run()["coord_update"] + shift = np.array([1.25, -0.5, 0.75]) + shifted_node, shifted_latent = self.desc.call_with_latent( + self.coord + shift, + self.atype, + self.nlist, + mapping=np.tile(np.arange(self.nloc), (1, 1)), + ) + shifted = self.fitting.call_atoms( + shifted_node, shifted_latent, self.nlist, self.coord + shift + )["coord_update"] + np.testing.assert_allclose( + shifted - base, + np.broadcast_to(shift, base.shape), + rtol=1e-10, + atol=1e-10, + ) + + def test_columns_past_the_frame_are_not_covered(self) -> None: + """The declared width is fixed; a shorter frame pads the rest.""" + out = self._run() + np.testing.assert_array_equal( + out["pair_mask"][:, :, self.nloc :], + np.zeros((1, self.nloc, self.max_atoms - self.nloc)), + ) + + def test_a_frame_wider_than_max_atoms_is_refused(self) -> None: + narrow = UniMolDPAPretrainFitting( + ntypes=3, + dim_descrpt=self.desc.channels, + node_readout_lmax=self.desc.node_readout_lmax, + max_atoms=self.nloc - 1, + dist_hidden=16, + precision="float64", + seed=1, + ) + node, latent = self.desc.call_with_latent( + self.coord, + self.atype, + self.nlist, + mapping=np.tile(np.arange(self.nloc), (1, 1)), + ) + with self.assertRaisesRegex(ValueError, "exceeds max_atoms"): + narrow.call_atoms(node, latent, self.nlist, self.coord) + + def test_the_five_tuple_path_is_refused(self) -> None: + """These heads need the equivariant state and the neighbour list.""" + with self.assertRaisesRegex(NotImplementedError, "call_atoms"): + self.fitting.call(np.zeros((1, 1, 1)), np.zeros((1, 1), dtype=np.int64)) + + def test_optional_heads_can_all_be_disabled(self) -> None: + fitting = UniMolDPAPretrainFitting( + ntypes=3, + dim_descrpt=self.desc.channels, + node_readout_lmax=self.desc.node_readout_lmax, + max_atoms=self.max_atoms, + mask_token_head=False, + coord_head=False, + dist_head=False, + precision="float64", + seed=1, + ) + node, latent = self.desc.call_with_latent( + self.coord, + self.atype, + self.nlist, + mapping=np.tile(np.arange(self.nloc), (1, 1)), + ) + self.assertEqual(fitting.output_def().var_defs, {}) + self.assertEqual(fitting.call_atoms(node, latent, self.nlist, self.coord), {}) + + def test_serialize_round_trip(self) -> None: + revived = UniMolDPAPretrainFitting.deserialize(self.fitting.serialize()) + node, latent = self.desc.call_with_latent( + self.coord, + self.atype, + self.nlist, + mapping=np.tile(np.arange(self.nloc), (1, 1)), + ) + a = self.fitting.call_atoms(node, latent, self.nlist, self.coord) + b = revived.call_atoms(node, latent, self.nlist, self.coord) + for key in a: + with self.subTest(output=key): + np.testing.assert_allclose(a[key], b[key]) + + +class TestUniMolDPAAtomicModel(unittest.TestCase): + """The objective wired onto a DPA backbone, through the model machinery.""" + + def setUp(self) -> None: + rng = np.random.default_rng(0) + self.nloc, nnei, ntypes = 6, 5, 3 + self.type_map = ["C", "N", "[MASK]"] + self.desc = DescrptDPA4( + rcut=6.0, + rcut_smth=5.0, + sel=nnei, + ntypes=ntypes, + seed=0, + precision="float64", + ) + self.fitting = UniMolDPAPretrainFitting( + ntypes=ntypes, + dim_descrpt=self.desc.channels, + node_readout_lmax=self.desc.node_readout_lmax, + max_atoms=self.nloc, + dist_hidden=16, + precision="float64", + seed=1, + type_map=self.type_map, + ) + self.model = DPUniMolDPAAtomicModel(self.desc, self.fitting, self.type_map) + self.coord = rng.normal(size=(1, self.nloc, 3)) * 1.5 + self.atype = rng.integers(0, ntypes, size=(1, self.nloc)) + self.nlist = np.stack( + [ + np.stack( + [ + np.array([j for j in range(self.nloc) if j != i][:nnei]) + for i in range(self.nloc) + ] + ) + ] + ) + self.mapping = np.tile(np.arange(self.nloc), (1, 1)) + + def _predict(self): + return self.model.forward_common_atomic( + self.coord, self.atype, self.nlist, mapping=self.mapping + ) + + def test_the_heads_reach_the_model_output(self) -> None: + out = self._predict() + for name in ("token_logits", "coord_update", "pair_dist", "pair_mask"): + with self.subTest(output=name): + self.assertIn(name, out) + # the base class supplies the real-atom mask the objective needs + self.assertIn("mask", out) + + def test_a_backbone_without_the_equivariant_seam_is_refused(self) -> None: + """The coordinate head reads a state most descriptors do not expose.""" + from deepmd.dpmodel.descriptor.se_e2_a import ( + DescrptSeA, + ) + + plain = DescrptSeA(rcut=6.0, rcut_smth=5.0, sel=[10, 10, 10]) + with self.assertRaisesRegex(TypeError, "call_with_latent"): + DPUniMolDPAAtomicModel(plain, self.fitting, self.type_map) + + def test_the_objective_scores_it(self) -> None: + """The point of the whole branch: this output feeds Uni-Mol's loss. + + The backbone has no virtual tokens, and Uni-Mol's two norm regularisers + constrain quantities belonging to its own transformer, so they carry no + weight here. + """ + rng = np.random.default_rng(1) + pred = self._predict() + token_target = np.zeros((1, self.nloc), dtype=np.int64) + token_target[0, [1, 3]] = [5, 6] + labels = { + "unimol_token_target": token_target, + "unimol_coord_target": rng.normal(size=(1, self.nloc, 3)), + } + loss = UniMolLoss( + x_norm_loss=0.0, delta_pair_repr_norm_loss=0.0, virtual_tokens=False + ) + total, terms = loss.call(1.0, 0, pred, labels) + self.assertEqual(sorted(terms), ["coord_loss", "dist_loss", "token_loss"]) + for name, value in terms.items(): + with self.subTest(term=name): + self.assertTrue(bool(np.isfinite(float(value)))) + self.assertTrue(bool(np.isfinite(float(total)))) + + +class TestUniMolDPAChargeSpin(unittest.TestCase): + """A backbone with a charge/spin condition hears it through this model.""" + + def setUp(self) -> None: + rng = np.random.default_rng(0) + self.nloc, nnei, self.ntypes = 6, 5, 3 + self.type_map = ["C", "N", "[MASK]"] + self.coord = rng.normal(size=(1, self.nloc, 3)) * 1.5 + self.atype = rng.integers(0, self.ntypes, size=(1, self.nloc)) + self.nlist = np.stack( + [ + np.stack( + [ + np.array([j for j in range(self.nloc) if j != i][:nnei]) + for i in range(self.nloc) + ] + ) + ] + ) + self.mapping = np.tile(np.arange(self.nloc), (1, 1)) + + def _model(self, **descriptor_kwargs): + desc = DescrptDPA4( + rcut=6.0, + rcut_smth=5.0, + sel=self.nlist.shape[-1], + ntypes=self.ntypes, + seed=0, + precision="float64", + **descriptor_kwargs, + ) + fitting = UniMolDPAPretrainFitting( + ntypes=self.ntypes, + dim_descrpt=desc.channels, + node_readout_lmax=desc.node_readout_lmax, + max_atoms=self.nloc, + dist_hidden=16, + precision="float64", + seed=1, + type_map=self.type_map, + ) + return DPUniMolDPAAtomicModel(desc, fitting, self.type_map) + + def _run(self, model, charge_spin=None): + return model.forward_common_atomic( + self.coord, + self.atype, + self.nlist, + mapping=self.mapping, + charge_spin=charge_spin, + ) + + def test_the_condition_reaches_the_outputs(self) -> None: + """The model advertises a charge/spin input, so it has to reach the heads. + + The coordinate head is left out: it reads the degree-one state, and an + untrained DPA4 cannot couple the scalar the condition enters into that + state yet, so it is unchanged until training moves the zero-initialised + mixing weights. It does not tell a dropped condition from a heard one. + """ + model = self._model(add_chg_spin_ebd=True, default_chg_spin=[0.0, 1.0]) + self.assertTrue(model.has_chg_spin_ebd()) + default = self._run(model) + other = self._run(model, np.array([[2.0, 3.0]])) + for name in ("token_logits", "pair_dist"): + with self.subTest(output=name): + self.assertGreater( + float(np.max(np.abs(np.asarray(default[name]) - other[name]))), + 1e-3, + ) + + def test_an_omitted_condition_means_the_default(self) -> None: + model = self._model(add_chg_spin_ebd=True, default_chg_spin=[0.0, 1.0]) + omitted = self._run(model) + explicit = self._run(model, np.array([[0.0, 1.0]])) + for name in ("token_logits", "coord_update", "pair_dist"): + with self.subTest(output=name): + np.testing.assert_allclose(omitted[name], explicit[name], atol=0) + + def test_a_backbone_without_the_condition_ignores_it(self) -> None: + model = self._model() + self.assertFalse(model.has_chg_spin_ebd()) + omitted = self._run(model) + given = self._run(model, np.array([[2.0, 3.0]])) + for name in ("token_logits", "coord_update", "pair_dist"): + with self.subTest(output=name): + np.testing.assert_allclose(omitted[name], given[name], atol=0) + + def test_a_condition_with_no_default_must_be_given(self) -> None: + model = self._model(add_chg_spin_ebd=True) + with self.assertRaisesRegex(ValueError, "charge_spin"): + self._run(model) + out = self._run(model, np.array([[2.0, 3.0]])) + self.assertTrue(bool(np.isfinite(np.asarray(out["token_logits"])).all())) + + +class TestUniMolDPAExample(unittest.TestCase): + """The shipped example must stay valid as the arguments evolve. + + Checked here rather than in the shared example test, because that one also + requires the referenced dataset to exist in the repository, and this example + points at data the user converts from upstream. + """ + + def test_example_configuration_is_valid(self) -> None: + import json + from pathlib import ( + Path, + ) + + from deepmd.utils.argcheck import ( + normalize, + ) + + path = ( + Path(__file__).parents[4] + / "examples" + / "unimol" + / "dpa_pretrain" + / "input.json" + ) + self.assertTrue(path.is_file(), path) + config = normalize(json.loads(path.read_text()), multi_task=False) + + fitting = config["model"]["fitting_net"] + self.assertEqual(fitting["type"], "unimol_dpa_pretrain") + self.assertEqual(fitting["dist_coverage"], "neighbour") + self.assertIn("[MASK]", config["model"]["type_map"]) + # a DPA backbone wraps nothing around a molecule, and Uni-Mol's two norm + # regularisers constrain its own transformer rather than this one + self.assertFalse(config["loss"]["virtual_tokens"]) + self.assertEqual(config["loss"]["x_norm_loss"], 0.0) + self.assertEqual(config["loss"]["delta_pair_repr_norm_loss"], 0.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt_expt/model/test_unimol_dpa.py b/source/tests/pt_expt/model/test_unimol_dpa.py new file mode 100644 index 0000000000..2e09960ae4 --- /dev/null +++ b/source/tests/pt_expt/model/test_unimol_dpa.py @@ -0,0 +1,354 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Uni-Mol's objective on a DPA backbone, through the configured path.""" + +import copy +import json +import os +import unittest + +import numpy as np +import torch + +from deepmd.pt_expt.loss.unimol import ( + UniMolLoss, +) +from deepmd.pt_expt.model import ( + get_model, +) +from deepmd.utils.argcheck import ( + normalize, +) + +HERE = os.path.dirname(os.path.abspath(__file__)) +EXAMPLE = os.path.join( + HERE, "..", "..", "..", "..", "examples", "water", "dpa4", "input.json" +) + + +def _config(coverage="neighbour", max_atoms=8): + """A small DPA4 backbone, built from the shipped example's own keys.""" + with open(EXAMPLE) as fh: + base = json.load(fh)["model"]["descriptor"] + desc = {k: v for k, v in base.items() if k != "_comment"} + desc.update( + { + "type": "dpa4", + "channels": 32, + "n_blocks": 1, + "precision": "float64", + "use_amp": False, + } + ) + return { + "model": { + "type_map": ["C", "N", "O", "H", "[MASK]"], + "descriptor": desc, + "fitting_net": { + "type": "unimol_dpa_pretrain", + "n_token": 31, + "max_atoms": max_atoms, + "dist_hidden": 16, + "dist_coverage": coverage, + "precision": "float64", + }, + }, + "learning_rate": {"type": "exp", "start_lr": 1e-4, "stop_lr": 1e-6}, + # the two norm regularisers constrain Uni-Mol's own transformer and have + # no counterpart on a DPA backbone, which also wraps nothing around a + # molecule + "loss": { + "type": "unimol", + "x_norm_loss": 0.0, + "delta_pair_repr_norm_loss": 0.0, + "virtual_tokens": False, + }, + "training": { + "training_data": {"systems": ["x"]}, + "numb_steps": 1, + "seed": 1, + }, + } + + +class TestUniMolDPAConfigured(unittest.TestCase): + def setUp(self) -> None: + self.nloc = 6 + rng = np.random.default_rng(0) + self.coord = rng.normal(size=(1, self.nloc, 3)) * 1.5 + self.atype = rng.integers(0, 4, size=(1, self.nloc)) + self.nlist = np.stack( + [ + np.stack( + [ + np.array([j for j in range(self.nloc) if j != i]) + for i in range(self.nloc) + ] + ) + ] + ) + + def _build(self, coverage="neighbour"): + return get_model(normalize(_config(coverage))["model"]) + + def _run(self, model): + device = next(model.parameters()).device + return model.forward_lower( + torch.as_tensor(self.coord, device=device), + torch.as_tensor(self.atype, device=device), + torch.as_tensor(self.nlist, device=device), + ) + + def test_the_configured_path_builds_and_runs(self) -> None: + model = self._build() + self.assertEqual(type(model).__name__, "UniMolDPAPretrainModel") + out = self._run(model) + for name in ("token_logits", "coord_update", "pair_dist", "pair_mask"): + with self.subTest(output=name): + self.assertIn(name, out) + self.assertTrue(bool(torch.isfinite(out[name]).all())) + + def test_a_partial_neighbour_list_covers_fewer_pairs(self) -> None: + """Neighbour coverage is a subset, which is the whole trade-off.""" + sparse = np.full((1, self.nloc, 2), -1, dtype=np.int64) + for i in range(self.nloc): + sparse[0, i] = [(i + 1) % self.nloc, (i + 2) % self.nloc] + full = self.nlist + + covered = {} + for name, nlist in (("neighbour", sparse), ("all_pairs", full)): + model = self._build("neighbour" if name == "neighbour" else "all_pairs") + device = next(model.parameters()).device + out = model.forward_lower( + torch.as_tensor(self.coord, device=device), + torch.as_tensor(self.atype, device=device), + torch.as_tensor(nlist, device=device), + ) + covered[name] = float(out["pair_mask"].sum().item()) + self.assertLess(covered["neighbour"], covered["all_pairs"]) + self.assertEqual(covered["neighbour"], 2 * self.nloc) + + def test_a_complete_neighbour_list_differs_only_by_the_diagonal(self) -> None: + """The two settings differ only by what a neighbour list cannot hold. + + Given a list that already holds every other atom, the only pairs left + over are the self-pairs -- which upstream scores and a neighbour list + can never contain. That is the whole of the difference, which is what + makes the option one of coverage rather than of model. + """ + covered = {} + for coverage in ("neighbour", "all_pairs"): + with self.subTest(coverage=coverage): + out = self._run(self._build(coverage)) + covered[coverage] = float(out["pair_mask"].sum().item()) + self.assertEqual(covered["all_pairs"] - covered["neighbour"], float(self.nloc)) + + def test_the_default_is_the_neighbour_list(self) -> None: + self.assertEqual( + normalize(_config())["model"]["fitting_net"]["dist_coverage"], "neighbour" + ) + + def test_the_coordinate_head_is_trainable(self) -> None: + """It projects with SO3Linear, whose weights are plain arrays. + + On this backend a plain array is registered as a buffer, and the + optimizer is built from the model's parameters, so without promotion + the head would train as a frozen random projection -- silently, because + the coordinate term still falls: gradients reach the backbone through + the frozen projection either way. + """ + model = self._build() + head = model.atomic_model.fitting_net.coord_head + names = [n for n, _ in head.named_parameters()] + self.assertIn("proj.weight", names) + self.assertIsInstance(head.proj.weight, torch.nn.Parameter) + self.assertTrue(head.proj.weight.requires_grad) + # and the optimizer would actually see it + self.assertTrue(any("coord_head" in n for n, _ in model.named_parameters())) + + def test_the_coordinate_head_actually_moves(self) -> None: + """Owning a parameter is not the same as learning one.""" + model = self._build() + head = model.atomic_model.fitting_net.coord_head + before = head.proj.weight.detach().clone() + + rng = np.random.default_rng(1) + device = next(model.parameters()).device + token_target = torch.zeros((1, self.nloc), dtype=torch.int64, device=device) + token_target[0, 1] = 5 + labels = { + "unimol_token_target": token_target, + "unimol_coord_target": torch.as_tensor( + rng.normal(size=(1, self.nloc, 3)), device=device + ), + } + params = copy.deepcopy(normalize(_config())["loss"]) + params.pop("type") + loss = UniMolLoss(**params) + opt = torch.optim.Adam(model.parameters(), lr=1e-2) + for _ in range(5): + opt.zero_grad() + total, _ = loss.call(1.0, 0, self._run(model), labels) + total.backward() + opt.step() + self.assertGreater(float((head.proj.weight.detach() - before).abs().max()), 0.0) + + def test_scoring_a_dpa_backbone_with_virtual_tokens_is_refused(self) -> None: + """Getting this wrong shifts every distance label by one column. + + The two conventions differ in which columns of the pair axis are real. + Scoring a backbone that wraps nothing as though it wrapped BOS and EOS + selects the same number of entries, so it neither crashes nor changes + shape -- every corrupted row simply trains against another atom's + distances. + """ + out = self._run(self._build()) + rng = np.random.default_rng(1) + device = next(self._build().parameters()).device + token_target = torch.zeros((1, self.nloc), dtype=torch.int64, device=device) + token_target[0, 1] = 5 + labels = { + "unimol_token_target": token_target, + "unimol_coord_target": torch.as_tensor( + rng.normal(size=(1, self.nloc, 3)), device=device + ), + } + wrong = UniMolLoss( + x_norm_loss=0.0, delta_pair_repr_norm_loss=0.0, virtual_tokens=True + ) + with self.assertRaisesRegex(ValueError, "wraps no virtual tokens"): + wrong.call(1.0, 0, out, labels) + + def test_a_periodic_frame_is_refused_on_the_path_evaluation_uses(self) -> None: + """The guard has to be where every path passes, not only at the top. + + A box handed to ``forward`` is the easy case. Evaluation and export go + through the lower path, where a cell has already become ghost atoms -- + which the coverage mask drops and the distance target has no + minimum-image convention for. Guarding only the upper path would leave + the case that matters unguarded. + """ + model = self._build() + device = next(model.parameters()).device + nall = self.nloc + 2 + coord = np.concatenate([self.coord, self.coord[:, :2] + 12.0], axis=1) + atype = np.concatenate([self.atype, self.atype[:, :2]], axis=1) + with self.assertRaisesRegex(ValueError, "ghost atom"): + model.forward_lower( + torch.as_tensor(coord, device=device), + torch.as_tensor(atype, device=device), + torch.as_tensor(self.nlist, device=device), + ) + self.assertEqual(nall, atype.shape[1]) + + def test_the_graph_path_is_refused_rather_than_half_advertised(self) -> None: + """The base class advertises a graph entry these heads do not have.""" + model = self._build() + self.assertFalse(model.atomic_model.supports_graph_export()) + self.assertFalse(model.atomic_model.uses_graph_lower()) + + def test_the_objective_scores_the_configured_model(self) -> None: + model = self._build() + out = self._run(model) + rng = np.random.default_rng(1) + device = next(model.parameters()).device + token_target = torch.zeros((1, self.nloc), dtype=torch.int64, device=device) + token_target[0, 1] = 5 + token_target[0, 3] = 6 + labels = { + "unimol_token_target": token_target, + "unimol_coord_target": torch.as_tensor( + rng.normal(size=(1, self.nloc, 3)), device=device + ), + } + params = copy.deepcopy(normalize(_config())["loss"]) + params.pop("type") + total, terms = UniMolLoss(**params).call(1.0, 0, out, labels) + self.assertEqual(sorted(terms), ["coord_loss", "dist_loss", "token_loss"]) + self.assertTrue(bool(torch.isfinite(torch.as_tensor(total)))) + + def test_keeping_the_loss_defaults_on_this_backbone_is_refused(self) -> None: + """The two norm weights default to Uni-Mol's 0.01, unasked for. + + A configuration that selects this fitting and leaves the loss block + alone therefore weights two regularisers the backbone does not + produce, which used to surface as ``KeyError: 'x_norm'`` on the first + batch -- after the data had been read and the statistics computed. + The trainer checks it while the model and the loss are wired together, + so this goes through ``get_loss``, the function that does the wiring, + rather than calling the check directly. + """ + from deepmd.pt_expt.train.training import ( + get_loss, + ) + + cfg = normalize( + _config() | {"loss": {"type": "unimol", "virtual_tokens": False}} + ) + # the weights nobody wrote + self.assertEqual(cfg["loss"]["x_norm_loss"], 0.01) + self.assertEqual(cfg["loss"]["delta_pair_repr_norm_loss"], 0.01) + model = get_model(copy.deepcopy(cfg["model"])) + with self.assertRaises(ValueError) as caught: + get_loss( + copy.deepcopy(cfg["loss"]), + 1e-4, + len(cfg["model"]["type_map"]), + model, + ) + message = str(caught.exception) + for expected in ( + "x_norm", + "delta_pair_norm", + "x_norm_loss", + "delta_pair_repr_norm_loss", + ): + with self.subTest(names=expected): + self.assertIn(expected, message) + + def test_a_loss_built_by_hand_is_refused_the_same_way(self) -> None: + """Not every caller goes through the trainer.""" + out = self._run(self._build()) + device = next(self._build().parameters()).device + token_target = torch.zeros((1, self.nloc), dtype=torch.int64, device=device) + token_target[0, 1] = 5 + labels = { + "unimol_token_target": token_target, + "unimol_coord_target": torch.zeros( + (1, self.nloc, 3), dtype=torch.float64, device=device + ), + } + with self.assertRaisesRegex(ValueError, "emits no x_norm"): + UniMolLoss(virtual_tokens=False).call(1.0, 0, out, labels) + + def test_the_shipped_configuration_is_not_refused(self) -> None: + """A guard that fires on the supported case is worse than none.""" + from deepmd.pt_expt.train.training import ( + get_loss, + ) + + with open( + os.path.join( + HERE, + "..", + "..", + "..", + "..", + "examples", + "unimol", + "dpa_pretrain", + "input.json", + ) + ) as fh: + shipped = { + k: v for k, v in json.load(fh)["loss"].items() if not k.startswith("_") + } + cfg = normalize(_config() | {"loss": shipped}) + model = get_model(copy.deepcopy(cfg["model"])) + loss = get_loss( + copy.deepcopy(cfg["loss"]), 1e-4, len(cfg["model"]["type_map"]), model + ) + self.assertEqual(loss.x_norm_loss, 0.0) + + +if __name__ == "__main__": + unittest.main()