fix(pt): calibrate DPA4 bias and support PyTorch 2.13 compile

deepmd/dpmodel/atomic_model/inter_potential.py → deepmd/dpmodel/atomic_model/inner_potential.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Analytical pair potentials for Zone bridging (backend-agnostic port of
-``deepmd.pt``'s ``InterPotential``). Lives in the atomic-model package:
+``deepmd.pt``'s ``InnerPotential``). Lives in the atomic-model package:
 the atomic layer owns per-atom energy assembly, where the ZBL term is
 injected on the graph route.
 """
@@ -63,7 +63,7 @@
 _A_BOHR = 0.5291772109  # Bohr radius in Å
 
 
-class InterPotential(NativeOP):
+class InnerPotential(NativeOP):
     """Analytical pair potential for Zone bridging.
 
     Supports the Ziegler-Biersack-Littmark (ZBL) screened nuclear repulsion
@@ -72,7 +72,7 @@ class InterPotential(NativeOP):
     contributes ``V_ZBL(r_ij) / 2`` to both atom i and atom j, avoiding
     double-counting from the symmetric neighbor list. Backend-agnostic
     (array-API) port of the reference implementation in
-    ``deepmd.pt.model.model.sezm_model.InterPotential``.
+    ``deepmd.pt.model.model.sezm_model.InnerPotential``.
 
     Parameters
     ----------
@@ -93,7 +93,7 @@ def __init__(self, type_map: list[str], mode: str = "zbl") -> None:
         super().__init__()
         mode = str(mode).upper()
         if mode != "ZBL":
-            raise ValueError(f"Unknown InterPotential mode: {mode}")
+            raise ValueError(f"Unknown InnerPotential mode: {mode}")
         self.mode = mode
         self.type_map = list(type_map)
         self.ntypes_real = len(type_map)
@@ -266,8 +266,8 @@ def call(
         return xp.astype(xp.reshape(atom_energy, (1, n_node, 1)), edge_vec.dtype)
 
 
-@BaseAtomicModel.register("inter_potential")
-class InterPotentialAtomicModel(BaseAtomicModel):
+@BaseAtomicModel.register("inner_potential")
+class InnerPotentialAtomicModel(BaseAtomicModel):
     """Analytical bridging pair potential as an ATOMIC MODEL.
 
     First-principles composition design: the analytical term maps local
@@ -302,7 +302,7 @@ def __init__(
         **kwargs: Any,
     ) -> None:
         super().__init__(type_map, **kwargs)
-        self.potential = InterPotential(type_map=list(type_map), mode=mode)
+        self.potential = InnerPotential(type_map=list(type_map), mode=mode)
         self.mode = self.potential.mode
         self.rcut = float(rcut)
         self.sel = (
@@ -317,7 +317,7 @@ def change_type_map(
         If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types.
 
         The generic base handles the public map and the stat/exclusion state;
-        the element lookup belongs to :class:`InterPotential`, so the update is
+        the element lookup belongs to :class:`InnerPotential`, so the update is
         delegated there rather than reimplemented here (review 3649295675 --
         without it the lookup keeps the ORIGINAL elements while ``atype``
         values mean new ones, and a longer new map raises ``IndexError``).
@@ -403,7 +403,7 @@ def forward_atomic(
     ) -> dict:
         """Dense route unsupported: the term rides the NeighborGraph route only."""
         raise NotImplementedError(
-            "InterPotentialAtomicModel rides the NeighborGraph route only; "
+            "InnerPotentialAtomicModel rides the NeighborGraph route only; "
             "the dense (nlist) route has no injection site for the term"
         )
 
@@ -455,7 +455,7 @@ def serialize(self) -> dict:
         data.update(
             {
                 "@class": "Model",
-                "type": "inter_potential",
+                "type": "inner_potential",
                 "@version": 1,
                 "mode": self.mode,
                 "rcut": self.rcut,
@@ -465,7 +465,7 @@ def serialize(self) -> dict:
         return data
 
     @classmethod
-    def deserialize(cls, data: dict) -> "InterPotentialAtomicModel":
+    def deserialize(cls, data: dict) -> "InnerPotentialAtomicModel":
         data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("@class", None)

deepmd/dpmodel/model/dp_linear_model.py

@@ -35,7 +35,7 @@ class LinearEnergyModel(DPModelCommon, DPLinearModel_):
     energies; on the NeighborGraph route every child consumes the same
     graph, so the summed energy differentiates through one shared edge
     backward. Used e.g. for analytical bridging compositions
-    (learned model + :class:`~deepmd.dpmodel.atomic_model.inter_potential.InterPotentialAtomicModel`).
+    (learned model + :class:`~deepmd.dpmodel.atomic_model.inner_potential.InnerPotentialAtomicModel`).
     """
 
     def __init__(

deepmd/dpmodel/model/model.py

@@ -61,7 +61,7 @@ def get_standard_model(data: dict) -> BaseModel:
     data = copy.deepcopy(data)
     # Analytical bridging (e.g. ZBL): the radii feed the DESCRIPTOR's
     # InnerClamp/BridgingSwitch (mirrors pt's builder); the method builds the
-    # atomic model's InterPotential below.
+    # atomic model's InnerPotential below.
     bridging_method = str(data.get("bridging_method", "none"))
     bridging_enabled = bridging_method.lower() not in ("none", "")
     if bridging_enabled:
@@ -77,8 +77,8 @@ def get_standard_model(data: dict) -> BaseModel:
     # Composition, not a flag (first-principles design): the analytical
     # bridging term is its own atomic model, summed with the learned one by the
     # existing linear composition machinery.
-    from deepmd.dpmodel.atomic_model.inter_potential import (
-        InterPotentialAtomicModel,
+    from deepmd.dpmodel.atomic_model.inner_potential import (
+        InnerPotentialAtomicModel,
     )
     from deepmd.dpmodel.atomic_model.linear_atomic_model import (
         LinearEnergyAtomicModel,
@@ -87,7 +87,7 @@ def get_standard_model(data: dict) -> BaseModel:
         LinearEnergyModel,
     )
 
-    zbl_atomic = InterPotentialAtomicModel(
+    zbl_atomic = InnerPotentialAtomicModel(
         type_map=data["type_map"],
         mode=bridging_method,
         rcut=descriptor.get_rcut(),
@@ -140,7 +140,7 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
     exclusions and the analytical-bridging composition -- so ``spin`` and
     ``bridging_method`` combine for free: the wrapper re-classes whatever
     atomic model came back, be it a single learned model or a
-    ``LinearEnergyAtomicModel`` over ``[learned, InterPotential]`` (the
+    ``LinearEnergyAtomicModel`` over ``[learned, InnerPotential]`` (the
     analytical child accepts and ignores ``spin``; the learned child consumes
     it).
 

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -232,6 +232,8 @@ def _make_wrapped_sampler(
         @functools.lru_cache
         def wrapped_sampler() -> list[dict]:
             sampled = sampled_func()
+            if not sampled:
+                return sampled
             if self.pair_excl is not None:
                 pair_exclude_types = self.pair_excl.get_exclude_types()
                 for sample in sampled:
@@ -662,7 +664,9 @@ def compute_fitting_input_stat(
         """
         pass
 
-    def _get_forward_wrapper_func(self) -> Callable[..., torch.Tensor]:
+    def _get_forward_wrapper_func(
+        self,
+    ) -> Callable[..., dict[str, torch.Tensor]]:
         """Get a forward wrapper of the atomic model for output bias calculation."""
 
         def model_forward(
@@ -672,7 +676,9 @@ def model_forward(
             fparam: torch.Tensor | None = None,
             aparam: torch.Tensor | None = None,
             charge_spin: torch.Tensor | None = None,
+            spin: torch.Tensor | None = None,
         ) -> dict[str, torch.Tensor]:
+            del spin
             with (
                 torch.no_grad()
             ):  # it's essential for pure torch forward function to use auto_batchsize

deepmd/pt/model/model/make_model.py

@@ -37,6 +37,9 @@
     extend_input_and_build_neighbor_list,
     nlist_distinguish_types,
 )
+from deepmd.pt.utils.stat import (
+    compute_output_stats,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -304,6 +307,48 @@ def get_out_bias(self) -> torch.Tensor:
         def set_out_bias(self, out_bias: torch.Tensor) -> None:
             self.atomic_model.set_out_bias(out_bias)
 
+        def predict_atomic_outputs_for_stat(
+            self,
+            coord: torch.Tensor,
+            atype: torch.Tensor,
+            box: torch.Tensor | None,
+            fparam: torch.Tensor | None = None,
+            aparam: torch.Tensor | None = None,
+            charge_spin: torch.Tensor | None = None,
+            spin: torch.Tensor | None = None,
+        ) -> dict[str, torch.Tensor]:
+            """Return atomic outputs through the standard atomic-model path."""
+            return self.atomic_model._get_forward_wrapper_func()(
+                coord,
+                atype,
+                box,
+                fparam=fparam,
+                aparam=aparam,
+                charge_spin=charge_spin,
+                spin=spin,
+            )
+
+        def _change_out_bias_with_model_forward(
+            self,
+            merged: Callable[[], list[dict]] | list[dict],
+            model_forward: Callable[..., dict[str, torch.Tensor]],
+        ) -> None:
+            """Fit a residual output-bias shift from a complete model predictor."""
+            atomic_model = self.atomic_model
+            delta_bias, out_std = compute_output_stats(
+                merged,
+                atomic_model.get_ntypes(),
+                keys=atomic_model.bias_keys,
+                model_forward=model_forward,
+                rcond=atomic_model.rcond,
+                preset_bias=atomic_model.preset_out_bias,
+                stats_distinguish_types=(
+                    atomic_model.get_compute_stats_distinguish_types()
+                ),
+                intensive=atomic_model.get_intensive(),
+            )
+            atomic_model._store_out_stat(delta_bias, out_std, add=True)
+
         def change_out_bias(
             self,
             merged: Any,
@@ -326,6 +371,12 @@ def change_out_bias(
                         and do least square on the errors to obtain the target shift as bias.
                 'set-by-statistic' : directly use the statistic output bias in the target dataset.
             """
+            if bias_adjust_mode == "change-by-statistic":
+                self._change_out_bias_with_model_forward(
+                    merged,
+                    self.predict_atomic_outputs_for_stat,
+                )
+                return
             self.atomic_model.change_out_bias(
                 merged,
                 bias_adjust_mode=bias_adjust_mode,

deepmd/pt/model/model/model.py

@@ -49,6 +49,48 @@ def compute_or_load_stat(
         """
         raise NotImplementedError
 
+    def predict_atomic_outputs_for_stat(
+        self,
+        coord: torch.Tensor,
+        atype: torch.Tensor,
+        box: torch.Tensor | None,
+        fparam: torch.Tensor | None = None,
+        aparam: torch.Tensor | None = None,
+        charge_spin: torch.Tensor | None = None,
+        spin: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """
+        Return complete atomic outputs used by residual output statistics.
+
+        Final model classes own this prediction contract because only they know
+        the complete physical forward, including model-level preprocessing and
+        analytical contributions. Implementations must not compute derivatives
+        or mutate compile caches.
+
+        Parameters
+        ----------
+        coord
+            Local coordinates with shape (nf, nloc, 3).
+        atype
+            Local atom types with shape (nf, nloc).
+        box
+            Simulation cells with shape (nf, 9), or ``None``.
+        fparam
+            Optional frame parameters.
+        aparam
+            Optional atomic parameters.
+        charge_spin
+            Optional frame-level charge and spin conditions.
+        spin
+            Optional native per-atom spin vectors.
+
+        Returns
+        -------
+        dict[str, torch.Tensor]
+            Complete atomic outputs for output-statistics regression.
+        """
+        raise NotImplementedError
+
     @torch.jit.export
     def get_observed_type_list(self) -> list[str]:
         """Get observed types (elements) of the model during data statistics.