add correction method to MultiSystems

dpdata/system.py

@@ -1294,6 +1294,45 @@ def pick_atom_idx(self, idx, nopbc=None):
             new_sys.append(ss.pick_atom_idx(idx, nopbc=nopbc))
         return new_sys
 
+    def correction(self, hl_sys: "MultiSystems"):
+        """Get energy and force correction between self (assumed low-level) and a high-level MultiSystems.
+        The self's coordinates will be kept, but energy and forces will be replaced by
+        the correction between these two systems.
+
+        Notes
+        -----
+        This method will not check whether coordinates and elements of two systems
+        are the same. The user should make sure by itself.
+
+        Parameters
+        ----------
+        hl_sys : MultiSystems
+            high-level MultiSystems
+
+        Returns
+        -------
+        corrected_sys : MultiSystems
+            Corrected MultiSystems
+
+        Examples
+        --------
+        Get correction between a low-level system and a high-level system:
+
+        >>> low_level = dpdata.MultiSystems().from_deepmd_hdf5("low_level.hdf5")
+        >>> high_level = dpdata.MultiSystems().from_deepmd_hdf5("high_level.hdf5")
+        >>> corr = low_level.correction(high_lebel)
+        >>> corr.to_deepmd_hdf5("corr.hdf5")
+        """
+        if not isinstance(hl_sys, MultiSystems):
+            raise RuntimeError("high_sys should be MultiSystems")
+        corrected_sys = MultiSystems(type_map=self.atom_names)
+        for nn in self.systems.keys():
+            ll_ss = self[nn]
+            hl_ss = hl_sys[nn]
+            corrected_sys.append(ll_ss.correction(hl_ss))
+        return corrected_sys
+
+
 def get_cls_name(cls: object) -> str:
     """Returns the fully qualified name of a class, such as `np.ndarray`.
 

tests/test_corr.py

@@ -28,5 +28,19 @@ def setUp(self):
         self.f_places = 6
         self.v_places = 6
 
+
+class TestCorr(unittest.TestCase, CompLabeledSys, IsPBC):
+    """Make a test to get a correction of two MultiSystems."""
+    def setUp(self):
+        s_ll = dpdata.MultiSystems(dpdata.LabeledSystem("amber/corr/dp_ll", fmt="deepmd/npy"))
+        s_hl = dpdata.MultiSystems(dpdata.LabeledSystem("amber/corr/dp_hl", fmt="deepmd/npy"))
+        self.system_1 = tuple(s_ll.correction(s_hl).systems.values())[0]
+        self.system_2 = dpdata.LabeledSystem("amber/corr/dp_corr" ,fmt="deepmd/npy")
+        self.places = 5
+        self.e_places = 4
+        self.f_places = 6
+        self.v_places = 6
+
+
 if __name__ == '__main__':
     unittest.main()