feat(pt_expt): full model and refact the module output names of dpmodel backend

deepmd/dpmodel/fitting/general_fitting.py

@@ -584,6 +584,7 @@ def _call_common(
                 )
 
         # calculate the prediction
+        results: dict[str, Array] = {}
         if not self.mixed_types:
             outs = xp.zeros(
                 [nf, nloc, net_dim_out],
@@ -622,4 +623,5 @@ def _call_common(
         exclude_mask = xp.astype(exclude_mask, xp.bool)
         # nf x nloc x nod
         outs = xp.where(exclude_mask[:, :, None], outs, xp.zeros_like(outs))
-        return {self.var_name: outs}
+        results[self.var_name] = outs
+        return results

deepmd/dpmodel/infer/deep_eval.py

@@ -358,9 +358,7 @@ def _eval_model(
 
         results = []
         for odef in request_defs:
-            # it seems not doing conversion
-            # dp_name = self._OUTDEF_DP2BACKEND[odef.name]
-            dp_name = odef.name
+            dp_name = self._OUTDEF_DP2BACKEND[odef.name]
             if dp_name in batch_output:
                 shape = self._get_output_shape(odef, nframes, natoms)
                 if batch_output[dp_name] is not None:

deepmd/dpmodel/model/dipole_model.py

@@ -3,6 +3,9 @@
     Any,
 )
 
+from deepmd.dpmodel.array_api import (
+    Array,
+)
 from deepmd.dpmodel.atomic_model import (
     DPDipoleAtomicModel,
 )
@@ -31,3 +34,90 @@
     ) -> None:
         DPModelCommon.__init__(self)
         DPDipoleModel_.__init__(self, *args, **kwargs)
+
+    def translated_output_def(self) -> dict[str, Any]:
+        out_def_data = self.model_output_def().get_data()
+        output_def = {
+            "dipole": out_def_data["dipole"],
+            "global_dipole": out_def_data["dipole_redu"],
+        }
+        if self.do_grad_r("dipole"):
+            output_def["force"] = out_def_data["dipole_derv_r"]
+            output_def["force"].squeeze(-2)
+        if self.do_grad_c("dipole"):
+            output_def["virial"] = out_def_data["dipole_derv_c_redu"]
+            output_def["virial"].squeeze(-2)
+            output_def["atom_virial"] = out_def_data["dipole_derv_c"]
+            output_def["atom_virial"].squeeze(-2)
+        if "mask" in out_def_data:
+            output_def["mask"] = out_def_data["mask"]
+        return output_def
+
+    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,
+        )
+        if self.get_fitting_net() is not None:
+            model_predict = {}
+            model_predict["dipole"] = model_ret["dipole"]
+            model_predict["global_dipole"] = model_ret["dipole_redu"]
+            if self.do_grad_r("dipole"):
+                model_predict["force"] = model_ret.get("dipole_derv_r")
+            if self.do_grad_c("dipole"):
+                model_predict["virial"] = model_ret.get("dipole_derv_c_redu")
+                if do_atomic_virial:
+                    model_predict["atom_virial"] = model_ret.get("dipole_derv_c")
+            if "mask" in model_ret:
+                model_predict["mask"] = model_ret["mask"]
+        else:
+            model_predict = model_ret
+            model_predict["updated_coord"] += coord
+        return model_predict
+
+    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,
+        )
+        if self.get_fitting_net() is not None:
+            model_predict = {}
+            model_predict["dipole"] = model_ret["dipole"]
+            model_predict["global_dipole"] = model_ret["dipole_redu"]
+            if self.do_grad_r("dipole"):
+                model_predict["extended_force"] = model_ret.get("dipole_derv_r")
+            if self.do_grad_c("dipole"):
+                model_predict["virial"] = model_ret.get("dipole_derv_c_redu")
+                if do_atomic_virial:
+                    model_predict["extended_virial"] = model_ret.get("dipole_derv_c")
+        else:
+            model_predict = model_ret
+        return model_predict
+
+    forward_lower = call_lower

deepmd/dpmodel/model/dos_model.py

@@ -3,6 +3,9 @@
     Any,
 )
 
+from deepmd.dpmodel.array_api import (
+    Array,
+)
 from deepmd.dpmodel.atomic_model import (
     DPDOSAtomicModel,
 )
@@ -31,3 +34,70 @@
     ) -> None:
         DPModelCommon.__init__(self)
         DPDOSModel_.__init__(self, *args, **kwargs)
+
+    def translated_output_def(self) -> dict[str, Any]:
+        out_def_data = self.model_output_def().get_data()
+        output_def = {
+            "atom_dos": out_def_data["dos"],
+            "dos": out_def_data["dos_redu"],
+        }
+        if "mask" in out_def_data:
+            output_def["mask"] = out_def_data["mask"]
+        return output_def
+
+    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,
+        )
+        if self.get_fitting_net() is not None:
+            model_predict = {}
+            model_predict["atom_dos"] = model_ret["dos"]
+            model_predict["dos"] = model_ret["dos_redu"]
+            if "mask" in model_ret:
+                model_predict["mask"] = model_ret["mask"]
+        else:
+            model_predict = model_ret
+            model_predict["updated_coord"] += coord
+        return model_predict
+
+    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,
+        )
+        if self.get_fitting_net() is not None:
+            model_predict = {}
+            model_predict["atom_dos"] = model_ret["dos"]
+            model_predict["dos"] = model_ret["dos_redu"]
+        else:
+            model_predict = model_ret
+        return model_predict
+
+    forward_lower = call_lower

deepmd/dpmodel/model/dp_model.py

@@ -48,3 +48,7 @@ def update_sel(
     def get_fitting_net(self) -> BaseFitting:
         """Get the fitting network."""
         return self.atomic_model.fitting
+
+    def get_descriptor(self) -> BaseDescriptor:
+        """Get the descriptor."""
+        return self.atomic_model.descriptor

deepmd/dpmodel/model/dp_zbl_model.py

@@ -3,6 +3,9 @@
     Any,
 )
 
+from deepmd.dpmodel.array_api import (
+    Array,
+)
 from deepmd.dpmodel.atomic_model.linear_atomic_model import (
     DPZBLLinearEnergyAtomicModel,
 )
@@ -34,6 +37,114 @@
     ) -> None:
         super().__init__(*args, **kwargs)
 
+    def translated_output_def(self) -> dict[str, Any]:
+        out_def_data = self.model_output_def().get_data()
+        output_def = {
+            "atom_energy": out_def_data["energy"],
+            "energy": out_def_data["energy_redu"],
+        }
+        if self.do_grad_r("energy"):
+            output_def["force"] = out_def_data["energy_derv_r"]
+            output_def["force"].squeeze(-2)
+        if self.do_grad_c("energy"):
+            output_def["virial"] = out_def_data["energy_derv_c_redu"]
+            output_def["virial"].squeeze(-2)
+            output_def["atom_virial"] = out_def_data["energy_derv_c"]
+            output_def["atom_virial"].squeeze(-2)
+        if "mask" in out_def_data:
+            output_def["mask"] = out_def_data["mask"]
+        return output_def
+
+    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(-2)
+                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
+
+    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.get("energy_derv_r") is not None:
+                model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2)
+            else:
+                model_predict["extended_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["extended_virial"] = derv_c.squeeze(-2)
+                else:
+                    model_predict["extended_virial"] = derv_c
+        else:
+            if model_ret.get("dforce") is not None:
+                model_predict["dforce"] = model_ret["dforce"]
+        return model_predict
+
+    forward_lower = call_lower
+
     @classmethod
     def update_sel(
         cls,