diff --git a/dpdata/format.py b/dpdata/format.py index ee495d60f..b96c374a3 100644 --- a/dpdata/format.py +++ b/dpdata/format.py @@ -106,11 +106,11 @@ def from_multi_systems(self, directory, **kwargs): list of filenames """ if self.MultiMode == self.MultiModes.Directory: - return [name for name in os.listdir(directory) if os.path.isdir(os.path.join(directory, name))] + return [os.path.join(directory, name) for name in os.listdir(directory) if os.path.isdir(os.path.join(directory, name))] raise NotImplementedError("%s doesn't support MultiSystems.from" %(self.__class__.__name__)) def to_multi_systems(self, formulas, directory, **kwargs): if self.MultiMode == self.MultiModes.Directory: return [os.path.join(directory, ff) for ff in formulas] raise NotImplementedError("%s doesn't support MultiSystems.to" %(self.__class__.__name__)) - \ No newline at end of file + diff --git a/dpdata/plugins/deepmd.py b/dpdata/plugins/deepmd.py index 6910cc60a..c2636e993 100644 --- a/dpdata/plugins/deepmd.py +++ b/dpdata/plugins/deepmd.py @@ -15,7 +15,7 @@ def to_system(self, data, file_name, **kwargs): """ dpdata.deepmd.raw.dump(file_name, data) - def from_labeled_system(self, file_name, type_map, **kwargs): + def from_labeled_system(self, file_name, type_map=None, **kwargs): return dpdata.deepmd.raw.to_system_data(file_name, type_map=type_map, labels=True) MultiMode = Format.MultiModes.Directory @@ -50,7 +50,7 @@ def to_system(self, data, file_name, set_size=5000, prec=np.float32, **kwargs): dpdata.deepmd.comp.dump( file_name, data, set_size=set_size, comp_prec=prec) - def from_labeled_system(self, file_name, type_map, **kwargs): + def from_labeled_system(self, file_name, type_map=None, **kwargs): return dpdata.deepmd.comp.to_system_data(file_name, type_map=type_map, labels=True) MultiMode = Format.MultiModes.Directory diff --git a/tests/test_multisystems.py b/tests/test_multisystems.py index df8d318b6..6e05861b3 100644 --- a/tests/test_multisystems.py +++ b/tests/test_multisystems.py @@ -30,6 +30,15 @@ def setUp(self): self.system_sizes = {'C1H4':2, 'C1H3':1} self.atom_names = ['C', 'H'] + def test_len(self): + self.assertEqual(len(self.systems), 2) + + def test_get_nframes(self): + self.assertEqual(self.systems.get_nframes(), 3) + + def test_str(self): + self.assertEqual(str(self.systems), "MultiSystems (2 systems containing 3 frames)") + class TestMultiSystemsAdd(unittest.TestCase, CompLabeledSys, MultiSystems, IsNoPBC): def setUp(self): @@ -117,5 +126,18 @@ def test_type_map(self): self.assertEqual(type_map, systems.atom_names) +class TestMultiSystemsTo(unittest.TestCase, MultiSystems): + def setUp(self): + # CH4 and O2 + system_1 = dpdata.LabeledSystem('gaussian/methane.gaussianlog', fmt='gaussian/log') + system_2 = dpdata.LabeledSystem('gaussian/oxygen.gaussianlog', fmt='gaussian/log') + systems1 = dpdata.MultiSystems(system_1, system_2) + systems1.to_deepmd_npy("tmp.multi.to") + self.systems = dpdata.MultiSystems().from_deepmd_npy("tmp.multi.to") + + self.system_names = ['C1H4O0', 'C0H0O2'] + self.system_sizes = {'C1H4O0':1, 'C0H0O2':1} + self.atom_names = ['C', 'H', 'O'] + if __name__ == '__main__': unittest.main()