fix: add more checks

deepmodeling · Jan 15, 2025 · f05f9ca · f05f9ca
1 parent 944af9d
commit f05f9ca
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/dpdata/driver.py b/dpdata/driver.py
@@ -166,7 +166,8 @@ def label(self, data: dict) -> dict:
                 labeled_data = lb_data.copy()
             else:
                 labeled_data["energies"] += lb_data["energies"]
-                labeled_data["forces"] += lb_data["forces"]
+                if "forces" in labeled_data and "forces" in lb_data:
+                    labeled_data["forces"] += lb_data["forces"]
                 if "virials" in labeled_data and "virials" in lb_data:
                     labeled_data["virials"] += lb_data["virials"]
         return labeled_data

diff --git a/dpdata/plugins/ase.py b/dpdata/plugins/ase.py
@@ -298,7 +298,8 @@ def from_labeled_system(
             dict_frames["energies"] = np.append(
                 dict_frames["energies"], tmp["energies"][0]
             )
-            dict_frames["forces"] = np.append(dict_frames["forces"], tmp["forces"][0])
+            if "forces" in tmp.keys() and "forces" in dict_frames.keys():
+                dict_frames["forces"] = np.append(dict_frames["forces"], tmp["forces"][0])
             if "virials" in tmp.keys() and "virials" in dict_frames.keys():
                 dict_frames["virials"] = np.append(
                     dict_frames["virials"], tmp["virials"][0]
@@ -307,7 +308,8 @@ def from_labeled_system(
         ## Correct the shape of numpy arrays
         dict_frames["cells"] = dict_frames["cells"].reshape(-1, 3, 3)
         dict_frames["coords"] = dict_frames["coords"].reshape(len(sub_traj), -1, 3)
-        dict_frames["forces"] = dict_frames["forces"].reshape(len(sub_traj), -1, 3)
+        if "forces" in dict_frames.keys():
+            dict_frames["forces"] = dict_frames["forces"].reshape(len(sub_traj), -1, 3)
         if "virials" in dict_frames.keys():
             dict_frames["virials"] = dict_frames["virials"].reshape(-1, 3, 3)