Refactor property fitting interface

deepmd/dpmodel/fitting/property_fitting.py

@@ -88,9 +88,8 @@ def __init__(
     ) -> None:
         self.task_dim = task_dim
         self.intensive = intensive
-        self.property_name = property_name
         super().__init__(
-            var_name=self.property_name,
+            var_name=property_name,
             ntypes=ntypes,
             dim_descrpt=dim_descrpt,
             dim_out=task_dim,
@@ -131,7 +130,7 @@ def serialize(self) -> dict:
             "type": "property",
             "task_dim": self.task_dim,
             "intensive": self.intensive,
-            "property_name": self.property_name,
+            "property_name": self.var_name,
         }
         dd["@version"] = 4
 

deepmd/entrypoints/test.py

@@ -779,12 +779,12 @@ def test_property(
     tuple[list[np.ndarray], list[int]]
         arrays with results and their shapes
     """
-    property_name = dp.get_property_name()
-    assert isinstance(property_name, str)
-    data.add(property_name, dp.task_dim, atomic=False, must=True, high_prec=True)
+    var_name = dp.get_var_name()
+    assert isinstance(var_name, str)
+    data.add(var_name, dp.task_dim, atomic=False, must=True, high_prec=True)
     if has_atom_property:
         data.add(
-            f"atom_{property_name}",
+            f"atom_{var_name}",
             dp.task_dim,
             atomic=True,
             must=False,
@@ -840,12 +840,12 @@ def test_property(
         aproperty = ret[1]
         aproperty = aproperty.reshape([numb_test, natoms * dp.task_dim])
 
-    diff_property = property - test_data[property_name][:numb_test]
+    diff_property = property - test_data[var_name][:numb_test]
     mae_property = mae(diff_property)
     rmse_property = rmse(diff_property)
 
     if has_atom_property:
-        diff_aproperty = aproperty - test_data[f"atom_{property_name}"][:numb_test]
+        diff_aproperty = aproperty - test_data[f"atom_{var_name}"][:numb_test]
         mae_aproperty = mae(diff_aproperty)
         rmse_aproperty = rmse(diff_aproperty)
 
@@ -862,7 +862,7 @@ def test_property(
         detail_path = Path(detail_file)
 
         for ii in range(numb_test):
-            test_out = test_data[property_name][ii].reshape(-1, 1)
+            test_out = test_data[var_name][ii].reshape(-1, 1)
             pred_out = property[ii].reshape(-1, 1)
 
             frame_output = np.hstack((test_out, pred_out))
@@ -876,7 +876,7 @@ def test_property(
 
         if has_atom_property:
             for ii in range(numb_test):
-                test_out = test_data[f"atom_{property_name}"][ii].reshape(-1, 1)
+                test_out = test_data[f"atom_{var_name}"][ii].reshape(-1, 1)
                 pred_out = aproperty[ii].reshape(-1, 1)
 
                 frame_output = np.hstack((test_out, pred_out))

deepmd/infer/deep_eval.py

@@ -274,9 +274,9 @@ def get_has_spin(self) -> bool:
         """Check if the model has spin atom types."""
         return False
 
-    def get_property_name(self) -> str:
+    def get_var_name(self) -> str:
         """Get the name of the fitting property."""
-        return NotImplementedError
+        raise NotImplementedError
 
     @abstractmethod
     def get_ntypes_spin(self) -> int:

deepmd/infer/deep_property.py

@@ -56,7 +56,7 @@ def change_output_def(self) -> None:
             FittingOutputDef(
                 [
                     OutputVariableDef(
-                        self.get_property_name(),
+                        self.get_var_name(),
                         shape=[self.get_task_dim()],
                         reducible=True,
                         atomic=True,
@@ -66,11 +66,11 @@ def change_output_def(self) -> None:
             )
         )
         self.deep_eval.output_def = self.output_def
-        self.deep_eval._OUTDEF_DP2BACKEND[self.get_property_name()] = (
-            f"atom_{self.get_property_name()}"
+        self.deep_eval._OUTDEF_DP2BACKEND[self.get_var_name()] = (
+            f"atom_{self.get_var_name()}"
         )
-        self.deep_eval._OUTDEF_DP2BACKEND[f"{self.get_property_name()}_redu"] = (
-            self.get_property_name()
+        self.deep_eval._OUTDEF_DP2BACKEND[f"{self.get_var_name()}_redu"] = (
+            self.get_var_name()
         )
 
     @property
@@ -136,10 +136,10 @@ def eval(
             aparam=aparam,
             **kwargs,
         )
-        atomic_property = results[self.get_property_name()].reshape(
+        atomic_property = results[self.get_var_name()].reshape(
             nframes, natoms, self.get_task_dim()
         )
-        property = results[f"{self.get_property_name()}_redu"].reshape(
+        property = results[f"{self.get_var_name()}_redu"].reshape(
             nframes, self.get_task_dim()
         )
 
@@ -159,9 +159,9 @@ def get_intensive(self) -> bool:
         """Get whether the property is intensive."""
         return self.deep_eval.get_intensive()
 
-    def get_property_name(self) -> str:
+    def get_var_name(self) -> str:
         """Get the name of the fitting property."""
-        return self.deep_eval.get_property_name()
+        return self.deep_eval.get_var_name()
 
 
 __all__ = ["DeepProperty"]

deepmd/pt/infer/deep_eval.py

@@ -184,9 +184,14 @@ def get_dim_aparam(self) -> int:
     def get_intensive(self) -> bool:
         return self.dp.model["Default"].get_intensive()
 
-    def get_property_name(self) -> str:
+    def get_var_name(self) -> str:
         """Get the name of the property."""
-        return self.dp.model["Default"].get_property_name()
+        if hasattr(self.dp.model["Default"], "get_var_name") and callable(
+            getattr(self.dp.model["Default"], "get_var_name")
+        ):
+            return self.dp.model["Default"].get_var_name()
+        else:
+            raise NotImplementedError
 
     @property
     def model_type(self) -> type["DeepEvalWrapper"]:
@@ -204,7 +209,7 @@ def model_type(self) -> type["DeepEvalWrapper"]:
             return DeepGlobalPolar
         elif "wfc" in model_output_type:
             return DeepWFC
-        elif self.dp.model["Default"].get_property_name() in model_output_type:
+        elif self.get_var_name() in model_output_type:
             return DeepProperty
         else:
             raise RuntimeError("Unknown model type")

deepmd/pt/loss/property.py

@@ -24,7 +24,7 @@ class PropertyLoss(TaskLoss):
     def __init__(
         self,
         task_dim,
-        property_name: str,
+        var_name: str,
         loss_func: str = "smooth_mae",
         metric: list = ["mae"],
         beta: float = 1.00,
@@ -39,6 +39,8 @@ def __init__(
         ----------
         task_dim : float
             The output dimension of property fitting net.
+        var_name : str
+            The atomic property to fit, 'energy', 'dipole', and 'polar'.
         loss_func : str
             The loss function, such as "smooth_mae", "mae", "rmse".
         metric : list
@@ -59,10 +61,10 @@ def __init__(
         self.loss_func = loss_func
         self.metric = metric
         self.beta = beta
-        self.property_name = property_name
         self.out_bias = out_bias
         self.out_std = out_std
         self.intensive = intensive
+        self.var_name = var_name
 
     def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False):
         """Return loss on properties .
@@ -88,12 +90,13 @@ def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False
             Other losses for display.
         """
         model_pred = model(**input_dict)
-        nbz = model_pred[self.property_name].shape[0]
-        assert model_pred[self.property_name].shape == (nbz, self.task_dim)
-        assert label[self.property_name].shape == (nbz, self.task_dim)
+        var_name = self.var_name
+        nbz = model_pred[var_name].shape[0]
+        assert model_pred[var_name].shape == (nbz, self.task_dim)
+        assert label[var_name].shape == (nbz, self.task_dim)
         if not self.intensive:
-            model_pred[self.property_name] = model_pred[self.property_name] / natoms
-            label[self.property_name] = label[self.property_name] / natoms
+            model_pred[var_name] = model_pred[var_name] / natoms
+            label[var_name] = label[var_name] / natoms
 
         if self.out_std is None:
             out_std = model.atomic_model.out_std[0][0]
@@ -123,28 +126,28 @@ def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False
         # loss
         if self.loss_func == "smooth_mae":
             loss += F.smooth_l1_loss(
-                (label[self.property_name] - out_bias) / out_std,
-                (model_pred[self.property_name] - out_bias) / out_std,
+                (label[var_name] - out_bias) / out_std,
+                (model_pred[var_name] - out_bias) / out_std,
                 reduction="sum",
                 beta=self.beta,
             )
         elif self.loss_func == "mae":
             loss += F.l1_loss(
-                (label[self.property_name] - out_bias) / out_std,
-                (model_pred[self.property_name] - out_bias) / out_std,
+                (label[var_name] - out_bias) / out_std,
+                (model_pred[var_name] - out_bias) / out_std,
                 reduction="sum",
             )
         elif self.loss_func == "mse":
             loss += F.mse_loss(
-                (label[self.property_name] - out_bias) / out_std,
-                (model_pred[self.property_name] - out_bias) / out_std,
+                (label[var_name] - out_bias) / out_std,
+                (model_pred[var_name] - out_bias) / out_std,
                 reduction="sum",
             )
         elif self.loss_func == "rmse":
             loss += torch.sqrt(
                 F.mse_loss(
-                    (label[self.property_name] - out_bias) / out_std,
-                    (model_pred[self.property_name] - out_bias) / out_std,
+                    (label[var_name] - out_bias) / out_std,
+                    (model_pred[var_name] - out_bias) / out_std,
                     reduction="mean",
                 )
             )
@@ -154,28 +157,28 @@ def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False
         # more loss
         if "smooth_mae" in self.metric:
             more_loss["smooth_mae"] = F.smooth_l1_loss(
-                label[self.property_name],
-                model_pred[self.property_name],
+                label[var_name],
+                model_pred[var_name],
                 reduction="mean",
                 beta=self.beta,
             ).detach()
         if "mae" in self.metric:
             more_loss["mae"] = F.l1_loss(
-                label[self.property_name],
-                model_pred[self.property_name],
+                label[var_name],
+                model_pred[var_name],
                 reduction="mean",
             ).detach()
         if "mse" in self.metric:
             more_loss["mse"] = F.mse_loss(
-                label[self.property_name],
-                model_pred[self.property_name],
+                label[var_name],
+                model_pred[var_name],
                 reduction="mean",
             ).detach()
         if "rmse" in self.metric:
             more_loss["rmse"] = torch.sqrt(
                 F.mse_loss(
-                    label[self.property_name],
-                    model_pred[self.property_name],
+                    label[var_name],
+                    model_pred[var_name],
                     reduction="mean",
                 )
             ).detach()
@@ -188,7 +191,7 @@ def label_requirement(self) -> list[DataRequirementItem]:
         label_requirement = []
         label_requirement.append(
             DataRequirementItem(
-                self.property_name,
+                self.var_name,
                 ndof=self.task_dim,
                 atomic=False,
                 must=True,

deepmd/pt/model/model/property_model.py

@@ -37,8 +37,8 @@ def __init__(
     def translated_output_def(self):
         out_def_data = self.model_output_def().get_data()
         output_def = {
-            f"atom_{self.get_property_name()}": out_def_data[self.get_property_name()],
-            self.get_property_name(): out_def_data[f"{self.get_property_name()}_redu"],
+            f"atom_{self.get_var_name()}": out_def_data[self.get_var_name()],
+            self.get_var_name(): out_def_data[f"{self.get_var_name()}_redu"],
         }
         if "mask" in out_def_data:
             output_def["mask"] = out_def_data["mask"]
@@ -62,12 +62,8 @@ def forward(
             do_atomic_virial=do_atomic_virial,
         )
         model_predict = {}
-        model_predict[f"atom_{self.get_property_name()}"] = model_ret[
-            self.get_property_name()
-        ]
-        model_predict[self.get_property_name()] = model_ret[
-            f"{self.get_property_name()}_redu"
-        ]
+        model_predict[f"atom_{self.get_var_name()}"] = model_ret[self.get_var_name()]
+        model_predict[self.get_var_name()] = model_ret[f"{self.get_var_name()}_redu"]
         if "mask" in model_ret:
             model_predict["mask"] = model_ret["mask"]
         return model_predict
@@ -80,12 +76,12 @@ def get_task_dim(self) -> int:
     @torch.jit.export
     def get_intensive(self) -> bool:
         """Get whether the property is intensive."""
-        return self.model_output_def()[self.get_property_name()].intensive
+        return self.model_output_def()[self.get_var_name()].intensive
 
     @torch.jit.export
-    def get_property_name(self) -> str:
+    def get_var_name(self) -> str:
         """Get the name of the property."""
-        return self.get_fitting_net().property_name
+        return self.get_fitting_net().var_name
 
     @torch.jit.export
     def forward_lower(
@@ -111,12 +107,8 @@ def forward_lower(
             extra_nlist_sort=self.need_sorted_nlist_for_lower(),
         )
         model_predict = {}
-        model_predict[f"atom_{self.get_property_name()}"] = model_ret[
-            self.get_property_name()
-        ]
-        model_predict[self.get_property_name()] = model_ret[
-            f"{self.get_property_name()}_redu"
-        ]
+        model_predict[f"atom_{self.get_var_name()}"] = model_ret[self.get_var_name()]
+        model_predict[self.get_var_name()] = model_ret[f"{self.get_var_name()}_redu"]
         if "mask" in model_ret:
             model_predict["mask"] = model_ret["mask"]
         return model_predict

deepmd/pt/model/task/property.py

@@ -93,9 +93,8 @@ def __init__(
     ) -> None:
         self.task_dim = task_dim
         self.intensive = intensive
-        self.property_name = property_name
         super().__init__(
-            var_name=self.property_name,
+            var_name=property_name,
             ntypes=ntypes,
             dim_descrpt=dim_descrpt,
             dim_out=task_dim,
@@ -147,7 +146,7 @@ def serialize(self) -> dict:
             "type": "property",
             "task_dim": self.task_dim,
             "intensive": self.intensive,
-            "property_name": self.property_name,
+            "property_name": self.var_name,
         }
         dd["@version"] = 4
 

deepmd/pt/train/training.py

@@ -1240,10 +1240,10 @@ def get_loss(loss_params, start_lr, _ntypes, _model):
         return TensorLoss(**loss_params)
     elif loss_type == "property":
         task_dim = _model.get_task_dim()
-        property_name = _model.get_property_name()
+        var_name = _model.get_var_name()
         intensive = _model.get_intensive()
         loss_params["task_dim"] = task_dim
-        loss_params["property_name"] = property_name
+        loss_params["var_name"] = var_name
         loss_params["intensive"] = intensive
         return PropertyLoss(**loss_params)
     else: