make output of energy model compatible among backends

Author: Han Wang

deepmd/dpmodel/model/ener_model.py

@@ -6,6 +6,9 @@
     Any,
 )
 
+from deepmd.dpmodel.array_api import (
+    Array,
+)
 from deepmd.dpmodel.atomic_model import (
     DPEnergyAtomicModel,
 )
@@ -48,6 +51,99 @@ def atomic_output_def(self) -> FittingOutputDef:
             return self.hess_fitting_def
         return super().atomic_output_def()
 
+    def call_lower(
+        self,
+        extended_coord: Array,
+        extended_atype: Array,
+        nlist: Array,
+        mapping: Array | None = None,
+        fparam: Array | None = None,
+        aparam: Array | None = None,
+        do_atomic_virial: bool = False,
+    ) -> dict[str, Array]:
+        model_ret = self.call_common_lower(
+            extended_coord,
+            extended_atype,
+            nlist,
+            mapping,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=do_atomic_virial,
+        )
+        model_predict = {}
+        model_predict["atom_energy"] = model_ret["energy"]
+        model_predict["energy"] = model_ret["energy_redu"]
+        if self.do_grad_r("energy"):
+            if model_ret["energy_derv_r"] is not None:
+                model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2)
+            else:
+                model_predict["extended_force"] = model_ret["energy_derv_r"]
+        if self.do_grad_c("energy"):
+            derv_c_redu = model_ret.get("energy_derv_c_redu")
+            if derv_c_redu is not None:
+                model_predict["virial"] = derv_c_redu.squeeze(-2)
+            else:
+                model_predict["virial"] = derv_c_redu
+            if do_atomic_virial:
+                if model_ret["energy_derv_c"] is not None:
+                    model_predict["extended_virial"] = model_ret[
+                        "energy_derv_c"
+                    ].squeeze(-3)
+                else:
+                    model_predict["extended_virial"] = model_ret["energy_derv_c"]
+        else:
+            if model_ret.get("dforce") is not None:
+                model_predict["dforce"] = model_ret["dforce"]
+        if "mask" in model_ret:
+            model_predict["mask"] = model_ret["mask"]
+        return model_predict
+
+    def call(
+        self,
+        coord: Array,
+        atype: Array,
+        box: Array | None = None,
+        fparam: Array | None = None,
+        aparam: Array | None = None,
+        do_atomic_virial: bool = False,
+    ) -> dict[str, Array]:
+        model_ret = self.call_common(
+            coord,
+            atype,
+            box,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=do_atomic_virial,
+        )
+        model_predict = {}
+        model_predict["atom_energy"] = model_ret["energy"]
+        model_predict["energy"] = model_ret["energy_redu"]
+        if self.do_grad_r("energy"):
+            if model_ret.get("energy_derv_r") is not None:
+                model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2)
+            else:
+                model_predict["force"] = model_ret.get("energy_derv_r")
+        if self.do_grad_c("energy"):
+            derv_c_redu = model_ret.get("energy_derv_c_redu")
+            if derv_c_redu is not None:
+                model_predict["virial"] = derv_c_redu.squeeze(-2)
+            else:
+                model_predict["virial"] = derv_c_redu
+            if do_atomic_virial:
+                derv_c = model_ret.get("energy_derv_c")
+                if derv_c is not None:
+                    model_predict["atom_virial"] = derv_c.squeeze(-3)
+                else:
+                    model_predict["atom_virial"] = derv_c
+        else:
+            if model_ret.get("dforce") is not None:
+                model_predict["force"] = model_ret["dforce"]
+        if "mask" in model_ret:
+            model_predict["mask"] = model_ret["mask"]
+        return model_predict
+
+    forward_lower = call_lower
+
     def translated_output_def(self) -> dict[str, Any]:
         """Get the translated output definition.
 

deepmd/dpmodel/model/make_model.py

@@ -223,7 +223,7 @@ def enable_compression(
                 check_frequency,
             )
 
-        def call(
+        def call_common(
             self,
             coord: Array,
             atype: Array,
@@ -232,7 +232,7 @@ def call(
             aparam: Array | None = None,
             do_atomic_virial: bool = False,
         ) -> dict[str, Array]:
-            """Return model prediction.
+            """Return model prediction with raw internal keys.
 
             Parameters
             ----------
@@ -262,7 +262,7 @@ def call(
             )
             del coord, box, fparam, aparam
             model_predict = model_call_from_call_lower(
-                call_lower=self.call_lower,
+                call_lower=self.call_common_lower,
                 rcut=self.get_rcut(),
                 sel=self.get_sel(),
                 mixed_types=self.mixed_types(),
@@ -277,7 +277,7 @@ def call(
             model_predict = self._output_type_cast(model_predict, input_prec)
             return model_predict
 
-        def call_lower(
+        def call_common_lower(
             self,
             extended_coord: Array,
             extended_atype: Array,
@@ -365,9 +365,10 @@ def forward_common_atomic(
                 mask=atomic_ret["mask"] if "mask" in atomic_ret else None,
             )
 
-        forward_lower = call_lower
-        forward_common = call
-        forward_common_lower = call_lower
+        call = call_common
+        call_lower = call_common_lower
+        forward_common = call_common
+        forward_common_lower = call_common_lower
 
         def get_out_bias(self) -> Array:
             """Get the output bias."""

deepmd/dpmodel/model/spin_model.py

@@ -377,7 +377,7 @@ def call(
         coord_updated, atype_updated = self.process_spin_input(coord, atype, spin)
         if aparam is not None:
             aparam = self.expand_aparam(aparam, nloc * 2)
-        model_predict = self.backbone_model.call(
+        model_predict = self.backbone_model.call_common(
             coord_updated,
             atype_updated,
             box,
@@ -447,7 +447,7 @@ def call_lower(
         )
         if aparam is not None:
             aparam = self.expand_aparam(aparam, nloc * 2)
-        model_predict = self.backbone_model.call_lower(
+        model_predict = self.backbone_model.call_common_lower(
             extended_coord_updated,
             extended_atype_updated,
             nlist_updated,

deepmd/pt_expt/model/ener_model.py

@@ -42,7 +42,7 @@ def forward(
         aparam: torch.Tensor | None = None,
         do_atomic_virial: bool = False,
     ) -> dict[str, torch.Tensor]:
-        model_ret = self.call(
+        model_ret = self.call_common(
             coord,
             atype,
             box,
@@ -73,7 +73,7 @@ def _forward_lower(
         aparam: torch.Tensor | None = None,
         do_atomic_virial: bool = False,
     ) -> dict[str, torch.Tensor]:
-        model_ret = self.call_lower(
+        model_ret = self.call_common_lower(
             extended_coord,
             extended_atype,
             nlist,
@@ -106,6 +106,26 @@ def forward_lower(
         fparam: torch.Tensor | None = None,
         aparam: torch.Tensor | None = None,
         do_atomic_virial: bool = False,
+    ) -> dict[str, torch.Tensor]:
+        return self._forward_lower(
+            extended_coord,
+            extended_atype,
+            nlist,
+            mapping,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=do_atomic_virial,
+        )
+
+    def forward_lower_exportable(
+        self,
+        extended_coord: torch.Tensor,
+        extended_atype: torch.Tensor,
+        nlist: torch.Tensor,
+        mapping: torch.Tensor | None = None,
+        fparam: torch.Tensor | None = None,
+        aparam: torch.Tensor | None = None,
+        do_atomic_virial: bool = False,
     ) -> torch.nn.Module:
         """Trace ``_forward_lower`` into an exportable module.
 
@@ -123,7 +143,7 @@ def forward_lower(
         torch.nn.Module
             A traced module whose ``forward`` accepts
             ``(extended_coord, extended_atype, nlist, mapping, fparam, aparam)``
-            and returns a dict with the same keys as ``_forward_lower``.
+            and returns a dict with the same keys as ``forward_lower``.
         """
         model = self
 

source/tests/consistent/model/test_ener.py

@@ -271,8 +271,8 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]:
         # shape not matched. ravel...
         if backend is self.RefBackend.DP:
             return (
-                ret["energy_redu"].ravel(),
                 ret["energy"].ravel(),
+                ret["atom_energy"].ravel(),
                 SKIP_FLAG,
                 SKIP_FLAG,
                 SKIP_FLAG,
@@ -303,11 +303,11 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]:
             )
         elif backend is self.RefBackend.JAX:
             return (
-                ret["energy_redu"].ravel(),
                 ret["energy"].ravel(),
-                ret["energy_derv_r"].ravel(),
-                ret["energy_derv_c_redu"].ravel(),
-                ret["energy_derv_c"].ravel(),
+                ret["atom_energy"].ravel(),
+                ret["force"].ravel(),
+                ret["virial"].ravel(),
+                ret["atom_virial"].ravel(),
             )
         elif backend is self.RefBackend.PD:
             return (
@@ -499,7 +499,7 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any:
         coord_tensor.requires_grad_(True)
         return {
             kk: vv.detach().cpu().numpy() if vv is not None else None
-            for kk, vv in pt_expt_obj.call_lower(
+            for kk, vv in pt_expt_obj.forward_lower(
                 coord_tensor,
                 pt_expt_numpy_to_torch(self.extended_atype),
                 pt_expt_numpy_to_torch(self.nlist),
@@ -536,8 +536,8 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]:
         # shape not matched. ravel...
         if backend is self.RefBackend.DP:
             return (
-                ret["energy_redu"].ravel(),
                 ret["energy"].ravel(),
+                ret["atom_energy"].ravel(),
                 SKIP_FLAG,
                 SKIP_FLAG,
                 SKIP_FLAG,
@@ -552,19 +552,19 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]:
             )
         elif backend is self.RefBackend.PT_EXPT:
             return (
-                ret["energy_redu"].ravel(),
                 ret["energy"].ravel(),
-                ret["energy_derv_r"].ravel(),
-                ret["energy_derv_c_redu"].ravel(),
-                ret["energy_derv_c"].ravel(),
+                ret["atom_energy"].ravel(),
+                ret["extended_force"].ravel(),
+                ret["virial"].ravel(),
+                ret["extended_virial"].ravel(),
             )
         elif backend is self.RefBackend.JAX:
             return (
-                ret["energy_redu"].ravel(),
                 ret["energy"].ravel(),
-                ret["energy_derv_r"].ravel(),
-                ret["energy_derv_c_redu"].ravel(),
-                ret["energy_derv_c"].ravel(),
+                ret["atom_energy"].ravel(),
+                ret["extended_force"].ravel(),
+                ret["virial"].ravel(),
+                ret["extended_virial"].ravel(),
             )
         elif backend is self.RefBackend.PD:
             return (
@@ -726,8 +726,8 @@ def test_set_out_bias(self) -> None:
         )
 
     def test_forward_common_alias(self) -> None:
-        """forward_common should be the same as call on dpmodel."""
-        ret_call = self.dp_model.call(
+        """forward_common should be the same as call_common on dpmodel."""
+        ret_call = self.dp_model.call_common(
             self.coords,
             self.atype,
             box=self.box,
@@ -741,8 +741,8 @@ def test_forward_common_alias(self) -> None:
             np.testing.assert_equal(ret_call[key], ret_fc[key])
 
     def test_forward_common_lower_alias(self) -> None:
-        """forward_common_lower should be the same as call_lower on dpmodel."""
-        ret_call = self.dp_model.call_lower(
+        """forward_common_lower should be the same as call_common_lower on dpmodel."""
+        ret_call = self.dp_model.call_common_lower(
             self.extended_coord,
             self.extended_atype,
             self.nlist,

source/tests/pt_expt/model/test_ener_model.py

@@ -143,11 +143,11 @@ def _prepare_lower_inputs(self):
         return ext_coord, ext_atype, nlist_t, mapping_t
 
     def test_forward_lower_exportable(self) -> None:
-        """Test that EnergyModel.forward_lower returns an exportable module.
+        """Test that EnergyModel.forward_lower_exportable returns an exportable module.
 
-        forward_lower() uses make_fx to trace through torch.autograd.grad,
-        decomposing the backward pass into primitive ops.  The returned module
-        can be passed directly to torch.export.export.
+        forward_lower_exportable() uses make_fx to trace through
+        torch.autograd.grad, decomposing the backward pass into primitive ops.
+        The returned module can be passed directly to torch.export.export.
 
         The test builds a model with numb_fparam > 0 and numb_aparam > 0 and
         verifies that:
@@ -184,7 +184,7 @@ def test_forward_lower_exportable(self) -> None:
         )
 
         # --- eager reference with zero params ---
-        ret_eager_zero = md._forward_lower(
+        ret_eager_zero = md.forward_lower(
             ext_coord.requires_grad_(True),
             ext_atype,
             nlist_t,
@@ -197,7 +197,7 @@ def test_forward_lower_exportable(self) -> None:
             self.assertIn(key, ret_eager_zero)
 
         # --- trace and export ---
-        traced = md.forward_lower(
+        traced = md.forward_lower_exportable(
             ext_coord,
             ext_atype,
             nlist_t,
@@ -262,7 +262,7 @@ def test_forward_lower_exportable(self) -> None:
             dtype=torch.float64,
             device=self.device,
         )
-        ret_eager_nz = md._forward_lower(
+        ret_eager_nz = md.forward_lower(
             ext_coord.requires_grad_(True),
             ext_atype,
             nlist_t,
@@ -364,13 +364,13 @@ def test_dp_consistency(self) -> None:
         ret_pt = md_pt(coord, self.atype, self.cell.reshape(1, 9))
 
         np.testing.assert_allclose(
-            ret_dp["energy_redu"],
+            ret_dp["energy"],
             ret_pt["energy"].detach().cpu().numpy(),
             rtol=1e-10,
             atol=1e-10,
         )
         np.testing.assert_allclose(
-            ret_dp["energy"],
+            ret_dp["atom_energy"],
             ret_pt["atom_energy"].detach().cpu().numpy(),
             rtol=1e-10,
             atol=1e-10,