Add split display in property fitting

Chengqian-Zhang · Nov 3, 2024 · bdbd631 · bdbd631
1 parent 0310ba6
commit bdbd631
Show file tree

Hide file tree

Showing 3 changed files with 88 additions and 48 deletions.
diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py
@@ -834,23 +834,29 @@ def test_property(
         aproperty = aproperty.reshape([numb_test, natoms * dp.task_dim])
 
     diff_property = property - test_data["property"][:numb_test]
+    mae_split_property = []
+    rmse_split_property = []
+
     mae_property = mae(diff_property)
     rmse_property = rmse(diff_property)
+    for jj in range(dp.task_dim):
+        mae_split_property.append(mae(diff_property[:, jj]))
+        rmse_split_property.append(rmse(diff_property[:, jj]))
 
     if has_atom_property:
-        diff_aproperty = aproperty - test_data["atom_property"][:numb_test]
-        mae_aproperty = mae(diff_aproperty)
-        rmse_aproperty = rmse(diff_aproperty)
+        #diff_aproperty = aproperty - test_data["atom_property"][:numb_test]
+        #mae_aproperty = mae(diff_aproperty)
+        #rmse_aproperty = rmse(diff_aproperty)
+        raise RuntimeError("Property dp test not support atom_property yet")
 
     log.info(f"# number of test data : {numb_test:d} ")
 
+    for jj in range(dp.task_dim):
+        log.info(f"PROPERTY MAE {jj:d} : {mae_split_property[jj]:e} units")
+        log.info(f"PROPERTY RMSE {jj:d} : {rmse_split_property[jj]:e} units")
     log.info(f"PROPERTY MAE            : {mae_property:e} units")
     log.info(f"PROPERTY RMSE           : {rmse_property:e} units")
 
-    if has_atom_property:
-        log.info(f"Atomic PROPERTY MAE     : {mae_aproperty:e} units")
-        log.info(f"Atomic PROPERTY RMSE    : {rmse_aproperty:e} units")
-
     if detail_file is not None:
         detail_path = Path(detail_file)
 
@@ -867,24 +873,14 @@ def test_property(
                 append=append_detail,
             )
 
-        if has_atom_property:
-            for ii in range(numb_test):
-                test_out = test_data["atom_property"][ii].reshape(-1, 1)
-                pred_out = aproperty[ii].reshape(-1, 1)
+    return_dict = {}
+    for jj in range(dp.task_dim):
+        return_dict[f"mae_property_{jj}"] = (mae_split_property[jj], property.shape[0])
+        return_dict[f"rmse_property_{jj}"] = (rmse_split_property[jj], property.shape[0])
+    return_dict["mae_property"] = (mae_property, property.size)
+    return_dict["rmse_property"] = (rmse_property, property.size)
 
-                frame_output = np.hstack((test_out, pred_out))
-
-                save_txt_file(
-                    detail_path.with_suffix(".aproperty.out.%.d" % ii),
-                    frame_output,
-                    header="%s - %.d: data_aproperty pred_aproperty" % (system, ii),
-                    append=append_detail,
-                )
-
-    return {
-        "mae_property": (mae_property, property.size),
-        "rmse_property": (rmse_property, property.size),
-    }
+    return return_dict
 
 
 def print_property_sys_avg(avg: dict[str, float]):
@@ -895,6 +891,10 @@ def print_property_sys_avg(avg: dict[str, float]):
     avg : np.ndarray
         array with summaries
     """
+    task_dim = int((len(avg.keys()) - 2)/2)
+    for jj in range(task_dim):
+        log.info(f"PROPERTY MAE {jj:d} : {avg[f'mae_property_{jj}']:e} units")
+        log.info(f"PROPERTY RMSE {jj:d} : {avg[f'rmse_property_{jj}']:e} units")
     log.info(f"PROPERTY MAE            : {avg['mae_property']:e} units")
     log.info(f"PROPERTY RMSE           : {avg['rmse_property']:e} units")
 

diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py
@@ -24,6 +24,7 @@ def __init__(
         loss_func: str = "smooth_mae",
         metric: list = ["mae"],
         beta: float = 1.00,
+        split_display: bool = False,
         **kwargs,
     ):
         r"""Construct a layer to compute loss on property.
@@ -44,6 +45,7 @@ def __init__(
         self.loss_func = loss_func
         self.metric = metric
         self.beta = beta
+        self.split_display = split_display
 
     def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False):
         """Return loss on properties .
@@ -104,33 +106,63 @@ def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False
             raise RuntimeError(f"Unknown loss function : {self.loss_func}")
 
         # more loss
-        if "smooth_mae" in self.metric:
-            more_loss["smooth_mae"] = F.smooth_l1_loss(
-                label["property"],
-                model_pred["property"],
-                reduction="mean",
-                beta=self.beta,
-            ).detach()
-        if "mae" in self.metric:
-            more_loss["mae"] = F.l1_loss(
-                label["property"],
-                model_pred["property"],
-                reduction="mean",
-            ).detach()
-        if "mse" in self.metric:
-            more_loss["mse"] = F.mse_loss(
-                label["property"],
-                model_pred["property"],
-                reduction="mean",
-            ).detach()
-        if "rmse" in self.metric:
-            more_loss["rmse"] = torch.sqrt(
-                F.mse_loss(
+        if self.split_display:
+            for jj in range(self.task_dim):
+                if "smooth_mae" in self.metric:
+                    more_loss[f"smooth_mae_{jj}"] = F.smooth_l1_loss(
+                        label["property"][:,jj],
+                        model_pred["property"][:,jj],
+                        reduction="mean",
+                        beta=self.beta,
+                    ).detach()
+                if "mae" in self.metric:
+                    more_loss[f"mae_{jj}"] = F.l1_loss(
+                        label["property"][:,jj],
+                        model_pred["property"][:,jj],
+                        reduction="mean",
+                    ).detach()
+                if "mse" in self.metric:
+                    more_loss[f"mse_{jj}"] = F.mse_loss(
+                        label["property"][:,jj],
+                        model_pred["property"][:,jj],
+                        reduction="mean",
+                    ).detach()
+                if "rmse" in self.metric:
+                    more_loss[f"rmse_{jj}"] = torch.sqrt(
+                        F.mse_loss(
+                            label["property"][:,jj],
+                            model_pred["property"][:,jj],
+                            reduction="mean",
+                        )
+                    ).detach()
+        else:
+            if "smooth_mae" in self.metric:
+                more_loss["smooth_mae"] = F.smooth_l1_loss(
                     label["property"],
                     model_pred["property"],
                     reduction="mean",
-                )
-            ).detach()
+                    beta=self.beta,
+                ).detach()
+            if "mae" in self.metric:
+                more_loss["mae"] = F.l1_loss(
+                    label["property"],
+                    model_pred["property"],
+                    reduction="mean",
+                ).detach()
+            if "mse" in self.metric:
+                more_loss["mse"] = F.mse_loss(
+                    label["property"],
+                    model_pred["property"],
+                    reduction="mean",
+                ).detach()
+            if "rmse" in self.metric:
+                more_loss["rmse"] = torch.sqrt(
+                    F.mse_loss(
+                        label["property"],
+                        model_pred["property"],
+                        reduction="mean",
+                    )
+                ).detach()
 
         return model_pred, loss, more_loss
 

diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py
@@ -2402,6 +2402,7 @@ def loss_property():
     doc_loss_func = "The loss function to minimize, such as 'mae','smooth_mae'."
     doc_metric = "The metric for display. This list can include 'smooth_mae', 'mae', 'mse' and 'rmse'."
     doc_beta = "The 'beta' parameter in 'smooth_mae' loss."
+    doc_split_display = "If display split property loss."
     return [
         Argument(
             "loss_func",
@@ -2424,6 +2425,13 @@ def loss_property():
             default=1.00,
             doc=doc_beta,
         ),
+        Argument(
+            "split_display",
+            bool,
+            optional=True,
+            default=False,
+            doc=doc_split_display,
+        ),
     ]