chore: refactor

deepmodeling · Jan 12, 2025 · 9c3c81e · 9c3c81e
1 parent 65efaa3
commit 9c3c81e
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/dpdata/system.py b/dpdata/system.py
@@ -1273,13 +1273,16 @@ def __add__(self, others):
             raise RuntimeError("Unspported data structure")
         return self.__class__.from_dict({"data": self_copy.data})
 
+    def has_forces(self) -> bool:
+        return "forces" in self.data
+
     def has_virial(self) -> bool:
         # return ('virials' in self.data) and (len(self.data['virials']) > 0)
         return "virials" in self.data
 
     def affine_map_fv(self, trans, f_idx: int | numbers.Integral):
         assert np.linalg.det(trans) != 0
-        if "forces" in self.data:
+        if self.has_forces():
             self.data["forces"][f_idx] = np.matmul(self.data["forces"][f_idx], trans)
         if self.has_virial():
             self.data["virials"][f_idx] = np.matmul(