From 15ec639b31094f56edc7ebc65e1ccdf4ace8554c Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 23 May 2022 00:27:06 -0400 Subject: [PATCH] refactor and optimize deepmd/hdf5 MultiSystems Now it is avoided to open the same file multiple times. --- dpdata/deepmd/hdf5.py | 14 +-- dpdata/plugins/deepmd.py | 173 ++++++++++++++++++++++++++++++++------ tests/test_deepmd_hdf5.py | 31 ++++++- 3 files changed, 185 insertions(+), 33 deletions(-) diff --git a/dpdata/deepmd/hdf5.py b/dpdata/deepmd/hdf5.py index 4b2486ba8..8bbd00484 100644 --- a/dpdata/deepmd/hdf5.py +++ b/dpdata/deepmd/hdf5.py @@ -1,4 +1,6 @@ """Utils for deepmd/hdf5 format.""" +from typing import Union + import h5py import numpy as np @@ -7,7 +9,7 @@ __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) : @@ -15,8 +17,8 @@ def to_system_data(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 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 diff --git a/dpdata/plugins/deepmd.py b/dpdata/plugins/deepmd.py index 2798d78ae..02ccf953b 100644 --- a/dpdata/plugins/deepmd.py +++ b/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") diff --git a/tests/test_deepmd_hdf5.py b/tests/test_deepmd_hdf5.py index 3c15b45e8..08d25730e 100644 --- a/tests/test_deepmd_hdf5.py +++ b/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')