feat(pt): add default_fparam

      "numb_fparam": 2,
      "default_fparam": [0.0, 1.0],

Author: iProzd

deepmd/entrypoints/test.py

@@ -298,7 +298,11 @@ def test_ener(
         data.add("atom_ener", 1, atomic=True, must=True, high_prec=False)
     if dp.get_dim_fparam() > 0:
         data.add(
-            "fparam", dp.get_dim_fparam(), atomic=False, must=True, high_prec=False
+            "fparam",
+            dp.get_dim_fparam(),
+            atomic=False,
+            must=not dp.has_default_fparam(),
+            high_prec=False,
         )
     if dp.get_dim_aparam() > 0:
         data.add("aparam", dp.get_dim_aparam(), atomic=True, must=True, high_prec=False)
@@ -334,7 +338,7 @@ def test_ener(
         atype = test_data["type"][:numb_test].reshape([numb_test, -1])
     else:
         atype = test_data["type"][0]
-    if dp.get_dim_fparam() > 0:
+    if dp.get_dim_fparam() > 0 and test_data["find_fparam"] != 0.0:
         fparam = test_data["fparam"][:numb_test]
     else:
         fparam = None

deepmd/infer/deep_eval.py

@@ -160,6 +160,9 @@ def get_type_map(self) -> list[str]:
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this DP."""
 
+    def has_default_fparam(self) -> bool:
+        return False
+
     @abstractmethod
     def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this DP."""
@@ -370,6 +373,9 @@ def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this DP."""
         return self.deep_eval.get_dim_fparam()
 
+    def has_default_fparam(self) -> bool:
+        return self.deep_eval.has_default_fparam()
+
     def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this DP."""
         return self.deep_eval.get_dim_aparam()

deepmd/pt/infer/deep_eval.py

@@ -183,6 +183,9 @@ def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this DP."""
         return self.dp.model["Default"].get_dim_aparam()
 
+    def has_default_fparam(self) -> bool:
+        return self.dp.model["Default"].has_default_fparam()
+
     def get_intensive(self) -> bool:
         return self.dp.model["Default"].get_intensive()
 

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -308,6 +308,9 @@ def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""
         return self.fitting_net.get_dim_fparam()
 
+    def has_default_fparam(self) -> bool:
+        return self.fitting_net.has_default_fparam()
+
     def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this atomic model."""
         return self.fitting_net.get_dim_aparam()

deepmd/pt/model/model/make_model.py

@@ -522,6 +522,10 @@ def get_dim_fparam(self) -> int:
             """Get the number (dimension) of frame parameters of this atomic model."""
             return self.atomic_model.get_dim_fparam()
 
+        @torch.jit.export
+        def has_default_fparam(self) -> bool:
+            return self.atomic_model.has_default_fparam()
+
         @torch.jit.export
         def get_dim_aparam(self) -> int:
             """Get the number (dimension) of atomic parameters of this atomic model."""

deepmd/pt/model/task/ener.py

@@ -56,6 +56,7 @@ def __init__(
         mixed_types: bool = True,
         seed: Optional[Union[int, list[int]]] = None,
         type_map: Optional[list[str]] = None,
+        default_fparam: Optional[list] = None,
         **kwargs,
     ) -> None:
         super().__init__(
@@ -74,6 +75,7 @@ def __init__(
             mixed_types=mixed_types,
             seed=seed,
             type_map=type_map,
+            default_fparam=default_fparam,
             **kwargs,
         )
 

deepmd/pt/model/task/fitting.py

@@ -227,6 +227,7 @@ def __init__(
         remove_vaccum_contribution: Optional[list[bool]] = None,
         type_map: Optional[list[str]] = None,
         use_aparam_as_mask: bool = False,
+        default_fparam: Optional[list] = None,
         **kwargs,
     ) -> None:
         super().__init__()
@@ -238,6 +239,7 @@ def __init__(
         self.resnet_dt = resnet_dt
         self.numb_fparam = numb_fparam
         self.numb_aparam = numb_aparam
+        self.default_fparam = default_fparam
         self.dim_case_embd = dim_case_embd
         self.activation_function = activation_function
         self.precision = precision
@@ -299,6 +301,20 @@ def __init__(
         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.register_buffer(
+                "default_fparam_tensor",
+                torch.tensor(
+                    np.array(self.default_fparam), dtype=self.prec, device=device
+                ),
+            )
+        else:
+            self.default_fparam_tensor = None
+
         in_dim = (
             self.dim_descrpt
             + self.numb_fparam
@@ -415,6 +431,9 @@ def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""
         return self.numb_fparam
 
+    def has_default_fparam(self) -> bool:
+        return self.default_fparam is not None
+
     def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this atomic model."""
         return self.numb_aparam
@@ -509,6 +528,13 @@ def _forward_common(
     ):
         # cast the input to internal precsion
         xx = descriptor.to(self.prec)
+        nf, nloc, nd = xx.shape
+
+        if self.numb_fparam > 0 and fparam is None:
+            # use default fparam
+            assert self.default_fparam_tensor is not None
+            fparam = torch.tile(self.default_fparam_tensor.unsqueeze(0), [nf, 1])
+
         fparam = fparam.to(self.prec) if fparam is not None else None
         aparam = aparam.to(self.prec) if aparam is not None else None
 
@@ -521,7 +547,6 @@ def _forward_common(
             xx_zeros = torch.zeros_like(xx)
         else:
             xx_zeros = None
-        nf, nloc, nd = xx.shape
         net_dim_out = self._net_out_dim()
 
         if nd != self.dim_descrpt:

deepmd/pt/model/task/invar_fitting.py

@@ -103,6 +103,7 @@ def __init__(
         atom_ener: Optional[list[Optional[torch.Tensor]]] = None,
         type_map: Optional[list[str]] = None,
         use_aparam_as_mask: bool = False,
+        default_fparam: Optional[list] = None,
         **kwargs,
     ) -> None:
         self.dim_out = dim_out
@@ -128,6 +129,7 @@ def __init__(
             else [x is not None for x in atom_ener],
             type_map=type_map,
             use_aparam_as_mask=use_aparam_as_mask,
+            default_fparam=default_fparam,
             **kwargs,
         )
 

deepmd/pt/train/training.py

@@ -1126,7 +1126,8 @@ def get_data(self, is_train=True, task_key="Default"):
         label_dict = {}
         for item_key in batch_data:
             if item_key in input_keys:
-                input_dict[item_key] = batch_data[item_key]
+                if item_key != "fparam" or batch_data["find_fparam"] != 0.0:
+                    input_dict[item_key] = batch_data[item_key]
             else:
                 if item_key not in ["sid", "fid"]:
                     label_dict[item_key] = batch_data[item_key]
@@ -1205,7 +1206,10 @@ def get_additional_data_requirement(_model):
     if _model.get_dim_fparam() > 0:
         fparam_requirement_items = [
             DataRequirementItem(
-                "fparam", _model.get_dim_fparam(), atomic=False, must=True
+                "fparam",
+                _model.get_dim_fparam(),
+                atomic=False,
+                must=not _model.has_default_fparam(),
             )
         ]
         additional_data_requirement += fparam_requirement_items

deepmd/utils/argcheck.py

@@ -1773,6 +1773,7 @@ def fitting_ener():
     return [
         Argument("numb_fparam", int, optional=True, default=0, doc=doc_numb_fparam),
         Argument("numb_aparam", int, optional=True, default=0, doc=doc_numb_aparam),
+        Argument("default_fparam", list, optional=True, default=None),
         Argument(
             "dim_case_embd",
             int,