fix virial in HybridDriver (#604)

Based on #603

---------

Co-authored-by: robinzyb <[email protected]>

Author: njzjz

dpdata/driver.py

@@ -163,6 +163,8 @@ def label(self, data: dict) -> dict:
             else:
                 labeled_data["energies"] += lb_data["energies"]
                 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
 
 

tests/comp_sys.py

@@ -92,8 +92,8 @@ def test_virial(self):
         # if len(self.system_1['virials']) == 0:
         #     self.assertEqual(len(self.system_1['virials']), 0)
         #     return
-        if "virials" not in self.system_1:
-            self.assertFalse("virials" in self.system_2)
+        if not self.system_1.has_virial():
+            self.assertFalse(self.system_2.has_virial())
             return
         np.testing.assert_almost_equal(
             self.system_1["virials"],

tests/test_predict.py

@@ -72,7 +72,7 @@ def setUp(self):
         self.system_2 = dpdata.LabeledSystem(
             "poscars/deepmd.h2o.md", fmt="deepmd/raw", type_map=["O", "H"]
         )
-        for pp in ("energies", "forces"):
+        for pp in ("energies", "forces", "virials"):
             self.system_2.data[pp][:] = 3.0
 
         self.places = 6