feat(dp/pt): add default_fparam (deepmodeling#4888)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Optional default frame parameter (default_fparam) across fitting nets
and model wrappers; new has_default_fparam() query exposed to users.

* **Behavior Changes**
* fparam is no longer always required when a default exists — data
loading, training, inference, and stats adapt to omit/use fparam;
requirements are conditional on has_default_fparam().
* TensorFlow paths accept/serialize default_fparam but disallow its
runtime use.

* **Documentation**
  * Docstrings and serializers updated; serialization versions bumped.

* **Tests**
  * Tests extended to cover explicit vs default fparam flows.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Duo <50307526+iProzd@users.noreply.github.com>

Author: iProzd

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -88,6 +88,10 @@ def get_type_map(self) -> list[str]:
         """Get the type map."""
         return self.type_map
 
+    def has_default_fparam(self) -> bool:
+        """Check if the model has default frame parameters."""
+        return False
+
     def reinit_atom_exclude(
         self,
         exclude_types: list[int] = [],

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -233,6 +233,10 @@ def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this atomic model."""
         return self.fitting.get_dim_aparam()
 
+    def has_default_fparam(self) -> bool:
+        """Check if the model has default frame parameters."""
+        return self.fitting.has_default_fparam()
+
     def get_sel_type(self) -> list[int]:
         """Get the selected atom types of this model.
 

deepmd/dpmodel/fitting/dipole_fitting.py

@@ -84,6 +84,9 @@ class DipoleFitting(GeneralFitting):
             Only reducible variable are differentiable.
     type_map: list[str], Optional
             A list of strings. Give the name to each type of atoms.
+    default_fparam: list[float], optional
+            The default frame parameter. If set, when `fparam.npy` files are not included in the data system,
+            this value will be used as the default value for the frame parameter in the fitting net.
     """
 
     def __init__(
@@ -110,6 +113,7 @@ def __init__(
         c_differentiable: bool = True,
         type_map: Optional[list[str]] = None,
         seed: Optional[Union[int, list[int]]] = None,
+        default_fparam: Optional[list[float]] = None,
     ) -> None:
         if tot_ener_zero:
             raise NotImplementedError("tot_ener_zero is not implemented")
@@ -144,6 +148,7 @@ def __init__(
             exclude_types=exclude_types,
             type_map=type_map,
             seed=seed,
+            default_fparam=default_fparam,
         )
 
     def _net_out_dim(self):
@@ -161,7 +166,7 @@ def serialize(self) -> dict:
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
         data = data.copy()
-        check_version_compatibility(data.pop("@version", 1), 3, 1)
+        check_version_compatibility(data.pop("@version", 1), 4, 1)
         var_name = data.pop("var_name", None)
         assert var_name == "dipole"
         return super().deserialize(data)

deepmd/dpmodel/fitting/dos_fitting.py

@@ -46,6 +46,7 @@ def __init__(
         exclude_types: list[int] = [],
         type_map: Optional[list[str]] = None,
         seed: Optional[Union[int, list[int]]] = None,
+        default_fparam: Optional[list] = None,
     ) -> None:
         if bias_dos is not None:
             self.bias_dos = bias_dos
@@ -70,12 +71,13 @@ def __init__(
             exclude_types=exclude_types,
             type_map=type_map,
             seed=seed,
+            default_fparam=default_fparam,
         )
 
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
         data = data.copy()
-        check_version_compatibility(data.pop("@version", 1), 3, 1)
+        check_version_compatibility(data.pop("@version", 1), 4, 1)
         data["numb_dos"] = data.pop("dim_out")
         data.pop("tot_ener_zero", None)
         data.pop("var_name", None)

deepmd/dpmodel/fitting/ener_fitting.py

@@ -46,6 +46,7 @@ def __init__(
         exclude_types: list[int] = [],
         type_map: Optional[list[str]] = None,
         seed: Optional[Union[int, list[int]]] = None,
+        default_fparam: Optional[list] = None,
     ) -> None:
         super().__init__(
             var_name="energy",
@@ -70,12 +71,13 @@ def __init__(
             exclude_types=exclude_types,
             type_map=type_map,
             seed=seed,
+            default_fparam=default_fparam,
         )
 
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
         data = data.copy()
-        check_version_compatibility(data.pop("@version", 1), 3, 1)
+        check_version_compatibility(data.pop("@version", 1), 4, 1)
         data.pop("var_name")
         data.pop("dim_out")
         return super().deserialize(data)

deepmd/dpmodel/fitting/general_fitting.py

@@ -94,6 +94,9 @@ class GeneralFitting(NativeOP, BaseFitting):
             A list of strings. Give the name to each type of atoms.
     seed: Optional[Union[int, list[int]]]
         Random seed for initializing the network parameters.
+    default_fparam: list[float], optional
+        The default frame parameter. If set, when `fparam.npy` files are not included in the data system,
+        this value will be used as the default value for the frame parameter in the fitting net.
     """
 
     def __init__(
@@ -120,6 +123,7 @@ def __init__(
         remove_vaccum_contribution: Optional[list[bool]] = None,
         type_map: Optional[list[str]] = None,
         seed: Optional[Union[int, list[int]]] = None,
+        default_fparam: Optional[list[float]] = None,
     ) -> None:
         self.var_name = var_name
         self.ntypes = ntypes
@@ -129,6 +133,7 @@ def __init__(
         self.numb_fparam = numb_fparam
         self.numb_aparam = numb_aparam
         self.dim_case_embd = dim_case_embd
+        self.default_fparam = default_fparam
         self.rcond = rcond
         self.tot_ener_zero = tot_ener_zero
         self.trainable = trainable
@@ -177,6 +182,15 @@ def __init__(
             self.case_embd = np.zeros(self.dim_case_embd, dtype=self.prec)
         else:
             self.case_embd = None
+
+        if self.default_fparam is not None:
+            if self.numb_fparam > 0:
+                assert len(self.default_fparam) == self.numb_fparam, (
+                    "default_fparam length mismatch!"
+                )
+            self.default_fparam_tensor = np.array(self.default_fparam, dtype=self.prec)
+        else:
+            self.default_fparam_tensor = None
         # init networks
         in_dim = (
             self.dim_descrpt
@@ -217,6 +231,10 @@ def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this atomic model."""
         return self.numb_aparam
 
+    def has_default_fparam(self) -> bool:
+        """Check if the fitting has default frame parameters."""
+        return self.default_fparam is not None
+
     def get_sel_type(self) -> list[int]:
         """Get the selected atom types of this model.
 
@@ -274,6 +292,8 @@ def __setitem__(self, key, value) -> None:
             self.case_embd = value
         elif key in ["scale"]:
             self.scale = value
+        elif key in ["default_fparam_tensor"]:
+            self.default_fparam_tensor = value
         else:
             raise KeyError(key)
 
@@ -292,6 +312,8 @@ def __getitem__(self, key):
             return self.case_embd
         elif key in ["scale"]:
             return self.scale
+        elif key in ["default_fparam_tensor"]:
+            return self.default_fparam_tensor
         else:
             raise KeyError(key)
 
@@ -306,7 +328,7 @@ def serialize(self) -> dict:
         """Serialize the fitting to dict."""
         return {
             "@class": "Fitting",
-            "@version": 3,
+            "@version": 4,
             "var_name": self.var_name,
             "ntypes": self.ntypes,
             "dim_descrpt": self.dim_descrpt,
@@ -315,6 +337,7 @@ def serialize(self) -> dict:
             "numb_fparam": self.numb_fparam,
             "numb_aparam": self.numb_aparam,
             "dim_case_embd": self.dim_case_embd,
+            "default_fparam": self.default_fparam,
             "rcond": self.rcond,
             "activation_function": self.activation_function,
             "precision": self.precision,
@@ -403,6 +426,14 @@ def _call_common(
             xx_zeros = xp.zeros_like(xx)
         else:
             xx_zeros = None
+
+        if self.numb_fparam > 0 and fparam is None:
+            # use default fparam
+            assert self.default_fparam_tensor is not None
+            fparam = xp.tile(
+                xp.reshape(self.default_fparam_tensor, (1, self.numb_fparam)), (nf, 1)
+            )
+
         # check fparam dim, concate to input descriptor
         if self.numb_fparam > 0:
             assert fparam is not None, "fparam should not be None"

deepmd/dpmodel/fitting/invar_fitting.py

@@ -110,6 +110,9 @@ class InvarFitting(GeneralFitting):
             Atomic contributions of the excluded atom types are set zero.
     type_map: list[str], Optional
             A list of strings. Give the name to each type of atoms.
+    default_fparam: list[float], optional
+        The default frame parameter. If set, when `fparam.npy` files are not included in the data system,
+        this value will be used as the default value for the frame parameter in the fitting net.
 
     """
 
@@ -138,6 +141,7 @@ def __init__(
         exclude_types: list[int] = [],
         type_map: Optional[list[str]] = None,
         seed: Optional[Union[int, list[int]]] = None,
+        default_fparam: Optional[list[float]] = None,
     ) -> None:
         if tot_ener_zero:
             raise NotImplementedError("tot_ener_zero is not implemented")
@@ -173,6 +177,7 @@ def __init__(
             else [x is not None for x in atom_ener],
             type_map=type_map,
             seed=seed,
+            default_fparam=default_fparam,
         )
 
     def serialize(self) -> dict:
@@ -185,7 +190,7 @@ def serialize(self) -> dict:
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
         data = data.copy()
-        check_version_compatibility(data.pop("@version", 1), 3, 1)
+        check_version_compatibility(data.pop("@version", 1), 4, 1)
         return super().deserialize(data)
 
     def _net_out_dim(self):

deepmd/dpmodel/fitting/polarizability_fitting.py

@@ -90,6 +90,9 @@ class PolarFitting(GeneralFitting):
             Whether to shift the diagonal part of the polarizability matrix. The shift operation is carried out after scale.
     type_map: list[str], Optional
             A list of strings. Give the name to each type of atoms.
+    default_fparam: list[float], optional
+            The default frame parameter. If set, when `fparam.npy` files are not included in the data system,
+            this value will be used as the default value for the frame parameter in the fitting net.
     """
 
     def __init__(
@@ -117,6 +120,7 @@ def __init__(
         shift_diag: bool = True,
         type_map: Optional[list[str]] = None,
         seed: Optional[Union[int, list[int]]] = None,
+        default_fparam: Optional[list[float]] = None,
     ) -> None:
         if tot_ener_zero:
             raise NotImplementedError("tot_ener_zero is not implemented")
@@ -164,6 +168,7 @@ def __init__(
             exclude_types=exclude_types,
             type_map=type_map,
             seed=seed,
+            default_fparam=default_fparam,
         )
 
     def _net_out_dim(self):
@@ -189,7 +194,7 @@ def __getitem__(self, key):
     def serialize(self) -> dict:
         data = super().serialize()
         data["type"] = "polar"
-        data["@version"] = 4
+        data["@version"] = 5
         data["embedding_width"] = self.embedding_width
         data["fit_diag"] = self.fit_diag
         data["shift_diag"] = self.shift_diag
@@ -200,7 +205,7 @@ def serialize(self) -> dict:
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
         data = data.copy()
-        check_version_compatibility(data.pop("@version", 1), 4, 1)
+        check_version_compatibility(data.pop("@version", 1), 5, 1)
         var_name = data.pop("var_name", None)
         assert var_name == "polar"
         return super().deserialize(data)

deepmd/dpmodel/fitting/property_fitting.py

@@ -65,6 +65,9 @@ class PropertyFittingNet(InvarFitting):
             Atomic contributions of the excluded atom types are set zero.
     type_map: list[str], Optional
             A list of strings. Give the name to each type of atoms.
+    default_fparam: list[float], optional
+            The default frame parameter. If set, when `fparam.npy` files are not included in the data system,
+            this value will be used as the default value for the frame parameter in the fitting net.
     """
 
     def __init__(
@@ -87,6 +90,7 @@ def __init__(
         mixed_types: bool = True,
         exclude_types: list[int] = [],
         type_map: Optional[list[str]] = None,
+        default_fparam: Optional[list] = None,
         # not used
         seed: Optional[int] = None,
     ) -> None:
@@ -110,6 +114,7 @@ def __init__(
             mixed_types=mixed_types,
             exclude_types=exclude_types,
             type_map=type_map,
+            default_fparam=default_fparam,
         )
 
     def output_def(self) -> FittingOutputDef:
@@ -129,7 +134,7 @@ def output_def(self) -> FittingOutputDef:
     @classmethod
     def deserialize(cls, data: dict) -> "PropertyFittingNet":
         data = data.copy()
-        check_version_compatibility(data.pop("@version"), 4, 1)
+        check_version_compatibility(data.pop("@version"), 5, 1)
         data.pop("dim_out")
         data["property_name"] = data.pop("var_name")
         data.pop("tot_ener_zero")
@@ -149,6 +154,6 @@ def serialize(self) -> dict:
             "task_dim": self.task_dim,
             "intensive": self.intensive,
         }
-        dd["@version"] = 4
+        dd["@version"] = 5
 
         return dd

deepmd/dpmodel/infer/deep_eval.py

@@ -120,6 +120,10 @@ def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this DP."""
         return self.dp.get_dim_aparam()
 
+    def has_default_fparam(self) -> bool:
+        """Check if the model has default frame parameters."""
+        return self.dp.has_default_fparam()
+
     @property
     def model_type(self) -> type["DeepEvalWrapper"]:
         """The the evaluator of the model type."""