feat(pt/dpmodel): add use_default_pf

Author: iProzd

deepmd/dpmodel/loss/ener.py

@@ -68,6 +68,9 @@ class EnergyLoss(Loss):
         The prefactor of generalized force loss at the end of the training.
     numb_generalized_coord : int
         The dimension of generalized coordinates.
+    use_default_pf : bool
+        If true, use default atom_pref of 1.0 for all atoms when atom_pref data is not provided.
+        This allows using the prefactor force loss (pf) without requiring atom_pref.npy files.
     use_huber : bool
         Enables Huber loss calculation for energy/force/virial terms with user-defined threshold delta (D).
         The loss function smoothly transitions between L2 and L1 loss:
@@ -110,6 +113,7 @@ def __init__(
         huber_delta: float = 0.01,
         loss_func: str = "mse",
         f_use_norm: bool = False,
+        use_default_pf: bool = False,
         **kwargs: Any,
     ) -> None:
         # Validate loss_func
@@ -149,6 +153,7 @@ def __init__(
         self.use_huber = use_huber
         self.huber_delta = huber_delta
         self.f_use_norm = f_use_norm
+        self.use_default_pf = use_default_pf
         if self.f_use_norm and not (self.use_huber or self.loss_func == "mae"):
             raise RuntimeError(
                 "f_use_norm can only be True when use_huber or loss_func='mae'."
@@ -182,7 +187,9 @@ def call(
         find_force = label_dict["find_force"]
         find_virial = label_dict["find_virial"]
         find_atom_ener = label_dict["find_atom_ener"]
-        find_atom_pref = label_dict["find_atom_pref"]
+        find_atom_pref = (
+            label_dict["find_atom_pref"] if not self.use_default_pf else 1.0
+        )
         xp = array_api_compat.array_namespace(
             energy,
             force,
@@ -477,6 +484,7 @@ def label_requirement(self) -> list[DataRequirementItem]:
                 must=False,
                 high_prec=False,
                 repeat=3,
+                default=1.0,
             )
         )
         if self.has_gf > 0:
@@ -512,7 +520,7 @@ def serialize(self) -> dict:
         """
         return {
             "@class": "EnergyLoss",
-            "@version": 2,
+            "@version": 3,
             "starter_learning_rate": self.starter_learning_rate,
             "start_pref_e": self.start_pref_e,
             "limit_pref_e": self.limit_pref_e,
@@ -533,6 +541,7 @@ def serialize(self) -> dict:
             "huber_delta": self.huber_delta,
             "loss_func": self.loss_func,
             "f_use_norm": self.f_use_norm,
+            "use_default_pf": self.use_default_pf,
         }
 
     @classmethod
@@ -550,6 +559,6 @@ def deserialize(cls, data: dict) -> "Loss":
             The deserialized loss module
         """
         data = data.copy()
-        check_version_compatibility(data.pop("@version"), 2, 1)
+        check_version_compatibility(data.pop("@version"), 3, 1)
         data.pop("@class")
         return cls(**data)

deepmd/pt/loss/ener.py

@@ -56,6 +56,7 @@ def __init__(
         loss_func: str = "mse",
         inference: bool = False,
         use_huber: bool = False,
+        use_default_pf: bool = False,
         f_use_norm: bool = False,
         huber_delta: float = 0.01,
         **kwargs: Any,
@@ -103,6 +104,9 @@ def __init__(
             MAE loss is less sensitive to outliers compared to MSE loss.
         inference : bool
             If true, it will output all losses found in output, ignoring the pre-factors.
+        use_default_pf : bool
+            If true, use default atom_pref of 1.0 for all atoms when atom_pref data is not provided.
+            This allows using the prefactor force loss (pf) without requiring atom_pref.npy files.
         use_huber : bool
             Enables Huber loss calculation for energy/force/virial terms with user-defined threshold delta (D).
             The loss function smoothly transitions between L2 and L1 loss:
@@ -147,6 +151,7 @@ def __init__(
         self.limit_pref_pf = limit_pref_pf
         self.start_pref_gf = start_pref_gf
         self.limit_pref_gf = limit_pref_gf
+        self.use_default_pf = use_default_pf
         self.relative_f = relative_f
         self.enable_atom_ener_coeff = enable_atom_ener_coeff
         self.numb_generalized_coord = numb_generalized_coord
@@ -357,7 +362,9 @@ def forward(
 
             if self.has_pf and "atom_pref" in label:
                 atom_pref = label["atom_pref"]
-                find_atom_pref = label.get("find_atom_pref", 0.0)
+                find_atom_pref = (
+                    label.get("find_atom_pref", 0.0) if not self.use_default_pf else 1.0
+                )
                 pref_pf = pref_pf * find_atom_pref
                 atom_pref_reshape = atom_pref.reshape(-1)
 
@@ -514,7 +521,7 @@ def label_requirement(self) -> list[DataRequirementItem]:
                     high_prec=True,
                 )
             )
-        if self.has_f:
+        if self.has_f or self.has_pf or self.relative_f is not None or self.has_gf:
             label_requirement.append(
                 DataRequirementItem(
                     "force",
@@ -553,6 +560,7 @@ def label_requirement(self) -> list[DataRequirementItem]:
                     must=False,
                     high_prec=False,
                     repeat=3,
+                    default=1.0,
                 )
             )
         if self.has_gf > 0:
@@ -588,7 +596,7 @@ def serialize(self) -> dict:
         """
         return {
             "@class": "EnergyLoss",
-            "@version": 2,
+            "@version": 3,
             "starter_learning_rate": self.starter_learning_rate,
             "start_pref_e": self.start_pref_e,
             "limit_pref_e": self.limit_pref_e,
@@ -609,6 +617,7 @@ def serialize(self) -> dict:
             "huber_delta": self.huber_delta,
             "loss_func": self.loss_func,
             "f_use_norm": self.f_use_norm,
+            "use_default_pf": self.use_default_pf,
         }
 
     @classmethod
@@ -626,7 +635,7 @@ def deserialize(cls, data: dict) -> "TaskLoss":
             The deserialized loss module
         """
         data = data.copy()
-        check_version_compatibility(data.pop("@version"), 2, 1)
+        check_version_compatibility(data.pop("@version"), 3, 1)
         data.pop("@class")
         return cls(**data)
 

deepmd/utils/argcheck.py

@@ -3189,6 +3189,11 @@ def loss_ener() -> list[Argument]:
         "atomic prefactor force", label="atom_pref", abbr="pf"
     )
     doc_limit_pref_pf = limit_pref("atomic prefactor force")
+    doc_use_default_pf = (
+        "If true, use default atom_pref of 1.0 for all atoms when atom_pref data is not provided. "
+        "This allows using the prefactor force loss (pf) without requiring atom_pref.npy files in training data. "
+        "When atom_pref.npy is provided, it will be used as-is regardless of this setting."
+    )
     doc_start_pref_gf = start_pref("generalized force", label="drdq", abbr="gf")
     doc_limit_pref_gf = limit_pref("generalized force")
     doc_numb_generalized_coord = "The dimension of generalized coordinates. Required when generalized force loss is used."
@@ -3299,6 +3304,13 @@ def loss_ener() -> list[Argument]:
             default=0.00,
             doc=doc_limit_pref_pf,
         ),
+        Argument(
+            "use_default_pf",
+            bool,
+            optional=True,
+            default=False,
+            doc=doc_use_default_pf,
+        ),
         Argument("relative_f", [float, None], optional=True, doc=doc_relative_f),
         Argument(
             "enable_atom_ener_coeff",

doc/model/train-se-a-mask.md

@@ -85,6 +85,22 @@ And the `loss` section in the training input script should be set as follows.
   }
 ```
 
+If `atom_pref.npy` is not provided in the training data, one can set `use_default_pf` to `true` to use a default atom preference of 1.0 for all atoms. This allows using the prefactor force loss (`pf` loss) without requiring `atom_pref.npy` files. When `atom_pref.npy` is provided, it will be used as-is regardless of this setting.
+
+```json
+"loss": {
+    "type": "ener",
+    "start_pref_e": 0.0,
+    "limit_pref_e": 0.0,
+    "start_pref_f": 0.0,
+    "limit_pref_f": 0.0,
+    "start_pref_pf": 1.0,
+    "limit_pref_pf": 1.0,
+    "use_default_pf": true,
+    "_comment": " that's all"
+  }
+```
+
 ## Type embedding
 
 Same as [`se_e2_a`](./train-se-e2-a.md).