register to funcs

Author: njzjz

README.md

@@ -131,26 +131,26 @@ Available properties are (nframe: number of frames in the system, natoms: total
 ## Dump data
 The data stored in `System` or `LabeledSystem` can be dumped in 'lammps/lmp' or 'vasp/poscar' format, for example:
 ```python
-d_outcar.to_lammps_lmp('conf.lmp', frame_idx=0)
+d_outcar.to('lammps/lmp', 'conf.lmp', frame_idx=0)
 ```
 The first frames of `d_outcar` will be dumped to 'conf.lmp'
 ```python
-d_outcar.to_vasp_poscar('POSCAR', frame_idx=-1)
+d_outcar.to('vasp/poscar', 'POSCAR', frame_idx=-1)
 ```
 The last frames of `d_outcar` will be dumped to 'POSCAR'.
 
 
 The data stored in `LabeledSystem` can be dumped to deepmd-kit raw format, for example
 ```python
-d_outcar.to_deepmd_raw('dpmd_raw')
+d_outcar.to('deepmd/raw', 'dpmd_raw')
 ```
 Or a simpler command:
 ```python
-dpdata.LabeledSystem('OUTCAR').to_deepmd_raw('dpmd_raw')
+dpdata.LabeledSystem('OUTCAR').to('deepmd/raw', 'dpmd_raw')
 ```
 Frame selection can be implemented by
 ```python
-dpdata.LabeledSystem('OUTCAR').sub_system([0,-1]).to_deepmd_raw('dpmd_raw')
+dpdata.LabeledSystem('OUTCAR').sub_system([0,-1]).to('deepmd/raw', 'dpmd_raw')
 ```
 by which only the first and last frames are dumped to `dpmd_raw`.
 

dpdata/system.py

@@ -116,8 +116,9 @@ def __init__ (self,
             self.apply_type_map(type_map)
 
     register_from_funcs = Register()
+    register_to_funcs = Register()
 
-    def from_fmt(self, file_name, fmt, **kwargs):
+    def from_fmt(self, file_name, fmt='auto', **kwargs):
         fmt = fmt.lower()
         if fmt == 'auto':
             fmt = os.path.basename(file_name).split('.')[-1].lower()
@@ -129,6 +130,17 @@ def from_fmt(self, file_name, fmt, **kwargs):
             func(self, file_name, **kwargs)
         else :
             raise RuntimeError('unknow data format ' + fmt)
+    
+    def to(self, fmt, *args, **kwargs):
+        fmt = fmt.lower()
+        to_funcs = self.register_to_funcs.funcs
+        if fmt in to_funcs:
+            func = to_funcs[fmt]
+            func_args = inspect.getfullargspec(func).args
+            kwargs = {kk: kwargs[kk] for kk in kwargs if kk in func_args}
+            func(self, *args, **kwargs)
+        else :
+            raise RuntimeError('unknow data format %s. Accepted format:' % (fmt, " ".join(to_funcs)))
 
     def __repr__(self):
         return self.__str__()
@@ -215,7 +227,7 @@ def map_atom_types(self,type_map=None):
 
         return new_atom_types
 
-
+    @register_to_funcs.register_funcs("list")
     def to_list(self):
         """
         convert system to list, usefull for data collection
@@ -430,6 +442,7 @@ def from_lammps_lmp (self, file_name, type_map = None) :
             self.data = dpdata.lammps.lmp.to_system_data(lines, type_map)
         self._shift_orig_zero()
 
+    @register_to_funcs.register_funcs("pymatgen/structure")
     def to_pymatgen_structure(self):
         '''
         convert System to Pymatgen Structure obj
@@ -449,6 +462,7 @@ def to_pymatgen_structure(self):
             structures.append(structure)
         return structures
 
+    @register_to_funcs.register_funcs("ase/structure")
     def to_ase_structure(self):
         '''
         convert System to ASE Atom obj
@@ -468,6 +482,7 @@ def to_ase_structure(self):
             structures.append(structure)
         return structures
 
+    @register_to_funcs.register_funcs("lammps/lmp")
     def to_lammps_lmp(self, file_name, frame_idx = 0) :
         """
         Dump the system in lammps data format
@@ -505,6 +520,7 @@ def from_vasp_poscar(self, file_name) :
             self.data = dpdata.vasp.poscar.to_system_data(lines)
         self.rot_lower_triangular()
 
+    @register_to_funcs.register_funcs("vasp/string")
     def to_vasp_string(self, frame_idx=0):
         """
         Dump the system in vasp POSCAR format string
@@ -518,6 +534,7 @@ def to_vasp_string(self, frame_idx=0):
         w_str = dpdata.vasp.poscar.from_system_data(self.data, frame_idx)
         return w_str
 
+    @register_to_funcs.register_funcs("vasp/poscar")
     def to_vasp_poscar(self, file_name, frame_idx = 0) :
         """
         Dump the system in vasp POSCAR format
@@ -556,6 +573,7 @@ def from_deepmd_raw(self, folder, type_map = None) :
         if tmp_data is not None :
             self.data = tmp_data
 
+    @register_to_funcs.register_funcs("deepmd/npy")
     def to_deepmd_npy(self, folder, set_size = 5000, prec=np.float32) :
         """
         Dump the system in deepmd compressed format (numpy binary) to `folder`.
@@ -578,6 +596,7 @@ def to_deepmd_npy(self, folder, set_size = 5000, prec=np.float32) :
                                 set_size = set_size,
                                 comp_prec = prec)
 
+    @register_to_funcs.register_funcs("deepmd/raw")
     def to_deepmd_raw(self, folder) :
         """
         Dump the system in deepmd raw format to `folder`

tests/test_lammps_lmp_dump.py

@@ -14,6 +14,15 @@ def setUp(self):
         self.system.from_lammps_lmp('tmp.lmp', 
                                     type_map = ['O', 'H'])
 
+class TestToFunc(unittest.TestCase, TestPOSCARoh):
+    
+    def setUp(self): 
+        tmp_system = dpdata.System(os.path.join('poscars', 'conf.lmp'), 
+                                   type_map = ['O', 'H'])
+        tmp_system.to('lammps/lmp', 'tmp.lmp')
+        self.system = dpdata.System()
+        self.system.from_fmt('tmp.lmp', fmt='lammps/lmp',
+                             type_map = ['O', 'H'])
 
 if __name__ == '__main__':
     unittest.main()