feat: run Uni-Mol's pretraining objective on a DPA4 backbone

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",

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

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(