refactor and optimize deepmd/hdf5 MultiSystems

dpdata/deepmd/hdf5.py

@@ -1,4 +1,6 @@
 """Utils for deepmd/hdf5 format."""
+from typing import Union
+
 import h5py
 import numpy as np
 
@@ -7,16 +9,16 @@
 
 __all__ = ['to_system_data', 'dump']
 
-def to_system_data(f: h5py.File,
+def to_system_data(f: Union[h5py.File, h5py.Group],
                    folder: str, 
                    type_map: list = None,
                    labels: bool = True) :
     """Load a HDF5 file.
 
     Parameters
     ----------
-    f : h5py.File
-        HDF5 file object
+    f : h5py.File or h5py.Group
+        HDF5 file or group object
     folder : str
         path in the HDF5 file
     type_map : list
@@ -82,7 +84,7 @@ def to_system_data(f: h5py.File,
         data['cells'] = np.zeros((nframes, 3, 3))
     return data
 
-def dump(f: h5py.File,
+def dump(f: Union[h5py.File, h5py.Group],
          folder: str,
          data: dict,
          set_size = 5000,
@@ -92,8 +94,8 @@ def dump(f: h5py.File,
 
     Parameters
     ----------
-    f : h5py.File
-        HDF5 file object
+    f : h5py.File or h5py.Group
+        HDF5 file or group object
     folder : str
         path in the HDF5 file
     data : dict

dpdata/plugins/deepmd.py

@@ -1,3 +1,5 @@
+from typing import Union, List
+
 import dpdata
 import dpdata.deepmd.raw
 import dpdata.deepmd.comp
@@ -69,41 +71,164 @@ class DeePMDHDF5Format(Format):
     >>> import dpdata
     >>> dpdata.MultiSystems().from_deepmd_npy("data").to_deepmd_hdf5("data.hdf5")
     """
-    def from_system(self, file_name, type_map=None, **kwargs):
-        s = file_name.split("#")
-        name = s[1] if len(s) > 1 else ""
-        with h5py.File(s[0], 'r') as f:
-            return dpdata.deepmd.hdf5.to_system_data(f, name, type_map=type_map, labels=False)
+    def _from_system(self, file_name: Union[str, h5py.Group, h5py.File], type_map: List[str], labels: bool):
+        """Convert HDF5 file to System or LabeledSystem data.
+
+        This method is used to switch from labeled or non-labeled options.
+
+        Parameters
+        ----------
+        file_name : str or h5py.Group or h5py.File
+            file name of the HDF5 file or HDF5 object. If it is a string,
+            hashtag is used to split path to the HDF5 file and the HDF5 group
+        type_map : dict[str]
+            type map
+        labels : bool
+            if Labeled
+
+        Returns
+        -------
+        dict
+            System or LabeledSystem data
+
+        Raises
+        ------
+        TypeError
+            file_name is not str or h5py.Group or h5py.File
+        """
+        if isinstance(file_name, (h5py.Group, h5py.File)):
+            return dpdata.deepmd.hdf5.to_system_data(file_name, "", type_map=type_map, labels=labels)
+        elif isinstance(file_name, str):
+            s = file_name.split("#")
+            name = s[1] if len(s) > 1 else ""
+            with h5py.File(s[0], 'r') as f:
+                return dpdata.deepmd.hdf5.to_system_data(f, name, type_map=type_map, labels=labels)
+        else:
+            raise TypeError("Unsupported file_name")
+
+    def from_system(self,
+                    file_name: Union[str, h5py.Group, h5py.File],
+                    type_map: List[str]=None,
+                    **kwargs) -> dict:
+        """Convert HDF5 file to System data.
+
+        Parameters
+        ----------
+        file_name : str or h5py.Group or h5py.File
+            file name of the HDF5 file or HDF5 object. If it is a string,
+            hashtag is used to split path to the HDF5 file and the HDF5 group
+        type_map : dict[str]
+            type map
+
+        Returns
+        -------
+        dict
+            System data
+
+        Raises
+        ------
+        TypeError
+            file_name is not str or h5py.Group or h5py.File
+        """
+        return self._from_system(file_name, type_map=type_map, labels=False)
+
+    def from_labeled_system(self,
+                            file_name: Union[str, h5py.Group, h5py.File],
+                            type_map: List[str]=None,
+                            **kwargs) -> dict:
+        """Convert HDF5 file to LabeledSystem data.
+
+        Parameters
+        ----------
+        file_name : str or h5py.Group or h5py.File
+            file name of the HDF5 file or HDF5 object. If it is a string,
+            hashtag is used to split path to the HDF5 file and the HDF5 group
+        type_map : dict[str]
+            type map
+
+        Returns
+        -------
+        dict
+            LabeledSystem data
+
+        Raises
+        ------
+        TypeError
+            file_name is not str or h5py.Group or h5py.File
+        """
+        return self._from_system(file_name, type_map=type_map, labels=True)
 
-    def from_labeled_system(self, file_name, type_map=None, **kwargs):
-        s = file_name.split("#")
-        name = s[1] if len(s) > 1 else ""
-        with h5py.File(s[0], 'r') as f:
-            return dpdata.deepmd.hdf5.to_system_data(f, name, type_map=type_map, labels=True)
-
     def to_system(self,
                   data : dict,
-                  file_name : str,
+                  file_name: Union[str, h5py.Group, h5py.File],
                   set_size : int = 5000, 
                   comp_prec : np.dtype = np.float64,
                   **kwargs):
-        s = file_name.split("#")
-        name = s[1] if len(s) > 1 else ""
-        mode = 'a' if name else 'w'
-        with h5py.File(s[0], mode) as f:
-            dpdata.deepmd.hdf5.dump(f, name, data, set_size = set_size, comp_prec = comp_prec)
+        """Convert System data to HDF5 file.
+
+        Parameters
+        ----------
+        data : dict
+            data dict
+        file_name : str or h5py.Group or h5py.File
+            file name of the HDF5 file or HDF5 object. If it is a string,
+            hashtag is used to split path to the HDF5 file and the HDF5 group
+        set_size : int, default=5000
+            set size
+        comp_prec : np.dtype
+            data precision
+        """
+        if isinstance(file_name, (h5py.Group, h5py.File)):
+            dpdata.deepmd.hdf5.dump(file_name, "", data, set_size = set_size, comp_prec = comp_prec)
+        elif isinstance(file_name, str):
+            s = file_name.split("#")
+            name = s[1] if len(s) > 1 else ""
+            with h5py.File(s[0], 'w') as f:
+                dpdata.deepmd.hdf5.dump(f, name, data, set_size = set_size, comp_prec = comp_prec)
+        else:
+            raise TypeError("Unsupported file_name")
 
     def from_multi_systems(self,
-                  directory,
-                  **kwargs):
+                  directory: str,
+                  **kwargs) -> h5py.Group:
+        """Generate HDF5 groups from a HDF5 file, which will be
+        passed to `from_system`.
+
+        Parameters
+        ----------
+        directory : str
+            HDF5 file name
+
+        Yields
+        ------
+        h5py.Group
+            a HDF5 group in the HDF5 file
+        """
         with h5py.File(directory, 'r') as f:
-            return ["%s#%s" % (directory, ff) for ff in f.keys()]
+            for ff in f.keys():
+                yield f[ff]
 
     def to_multi_systems(self,
-                  formulas,
-                  directory,
-                  **kwargs):
-        return ["%s#%s" % (directory, ff) for ff in formulas]
+                  formulas: List[str],
+                  directory: str,
+                  **kwargs) -> h5py.Group:
+        """Generate HDF5 groups, which will be passed to `to_system`.
+
+        Parameters
+        ----------
+        formulas : list[str]
+            formulas of MultiSystems
+        directory : str
+            HDF5 file name
+
+        Yields
+        ------
+        h5py.Group
+            a HDF5 group with the name of formula
+        """
+        with h5py.File(directory, 'w') as f:
+            for ff in formulas:
+                yield f.create_group(ff)
 
 
 @Driver.register("dp")

tests/test_deepmd_hdf5.py

@@ -2,9 +2,9 @@
 import numpy as np
 import unittest
 from context import dpdata
-from comp_sys import CompLabeledSys, CompSys, IsPBC
+from comp_sys import CompLabeledSys, CompSys, IsNoPBC, IsPBC, MultiSystems
 
-class TestDeepmdLoadDumpComp(unittest.TestCase, CompLabeledSys, IsPBC):
+class TestDeepmdLoadDumpHDF5(unittest.TestCase, CompLabeledSys, IsPBC):
     def setUp (self) :
         self.system_1 = dpdata.LabeledSystem('poscars/OUTCAR.h2o.md', 
                                              fmt = 'vasp/outcar')
@@ -25,7 +25,7 @@ def tearDown(self) :
             os.remove('tmp.deepmd.hdf5')
 
 
-class TestDeepmdCompNoLabels(unittest.TestCase, CompSys, IsPBC) :
+class TestDeepmdHDF5NoLabels(unittest.TestCase, CompSys, IsPBC) :
     def setUp (self) :
         self.system_1 = dpdata.System('poscars/POSCAR.h2o.md',
                                       fmt = 'vasp/poscar')
@@ -43,3 +43,28 @@ def setUp (self) :
     def tearDown(self) :
         if os.path.exists('tmp.deepmd.hdf5'):
             os.remove('tmp.deepmd.hdf5')
+
+
+class TestHDF5Multi(unittest.TestCase, CompLabeledSys, MultiSystems, IsNoPBC):
+    def setUp (self):
+        self.places = 6
+        self.e_places = 6
+        self.f_places = 6
+        self.v_places = 6
+
+        system_1 = dpdata.LabeledSystem('gaussian/methane.gaussianlog', fmt='gaussian/log')
+        system_2 = dpdata.LabeledSystem('gaussian/methane_reordered.gaussianlog', fmt='gaussian/log')
+        system_3 = dpdata.LabeledSystem('gaussian/methane_sub.gaussianlog', fmt='gaussian/log')
+        systems = dpdata.MultiSystems(system_1, system_2, system_3)
+        systems.to_deepmd_hdf5("tmp.deepmd.hdf5")
+
+        self.systems = dpdata.MultiSystems().from_deepmd_hdf5("tmp.deepmd.hdf5")
+        self.system_names = ['C1H4', 'C1H3']
+        self.system_sizes = {'C1H4':2, 'C1H3':1}
+        self.atom_names = ['C', 'H']
+        self.system_1 = self.systems['C1H3']
+        self.system_2 = system_3
+
+    def tearDown(self) :
+        if os.path.exists('tmp.deepmd.hdf5'):
+            os.remove('tmp.deepmd.hdf5')