Merge pull request #77 from njzjz/func_register

add experimental function register

Author: amcadmus

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

@@ -1,6 +1,7 @@
 #%%
 import os
 import glob
+import inspect
 import numpy as np
 import dpdata.lammps.lmp
 import dpdata.lammps.dump
@@ -23,6 +24,18 @@
 from dpdata.periodic_table import Element
 from dpdata.xyz.quip_gap_xyz import QuipGapxyzSystems
 
+
+class Register:
+    def __init__(self):
+        self.funcs = {}
+
+    def register_funcs(self, fmt):
+        def decorator(func):
+            self.funcs[fmt] = func
+            return func
+        return decorator
+
+
 class System (MSONable) :
     '''
     The data System
@@ -97,30 +110,37 @@ def __init__ (self,
             return
         if file_name is None :
             return
-        if fmt == 'auto':
-            fmt = os.path.basename(file_name).split('.')[-1]
-        if fmt == 'lmp' or fmt == 'lammps/lmp' :
-            self.from_lammps_lmp(file_name, type_map = type_map)
-        elif fmt == 'dump' or fmt == 'lammps/dump' :
-            self.from_lammps_dump(file_name, type_map = type_map, begin = begin, step = step)
-        elif fmt.lower() == 'poscar' or fmt.lower() == 'contcar' or fmt.lower() == 'vasp/poscar' or fmt.lower() == 'vasp/contcar':
-            self.from_vasp_poscar(file_name)
-        elif fmt == 'deepmd' or fmt == 'deepmd/raw':
-            self.from_deepmd_raw(file_name, type_map = type_map)
-        elif fmt == 'deepmd/npy':
-            self.from_deepmd_comp(file_name, type_map = type_map)
-        elif fmt == 'qe/cp/traj':
-            self.from_qe_cp_traj(file_name, begin = begin, step = step)
-        elif fmt.lower() == 'siesta/output':
-            self.from_siesta_output(file_name)
-        elif fmt.lower() == 'siesta/aimd_output':
-            self.from_siesta_aiMD_output(file_name)
-        else :
-            raise RuntimeError('unknow data format ' + fmt)
+        self.from_fmt(file_name, fmt, type_map=type_map, begin= begin, step=step)
 
         if type_map is not None:
             self.apply_type_map(type_map)
 
+    register_from_funcs = Register()
+    register_to_funcs = Register()
+
+    def from_fmt(self, file_name, fmt='auto', **kwargs):
+        fmt = fmt.lower()
+        if fmt == 'auto':
+            fmt = os.path.basename(file_name).split('.')[-1].lower()
+        from_funcs = self.register_from_funcs.funcs
+        if fmt in from_funcs:
+            func = from_funcs[fmt]
+            args = inspect.getfullargspec(func).args
+            kwargs = {kk: kwargs[kk] for kk in kwargs if kk in args}
+            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__()
@@ -207,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
@@ -414,13 +434,15 @@ def apply_pbc(self) :
         ncoord = ncoord % 1
         self.data['coords'] = np.matmul(ncoord, self.data['cells'])
 
-
+    @register_from_funcs.register_funcs("lmp")
+    @register_from_funcs.register_funcs("lammps/lmp")
     def from_lammps_lmp (self, file_name, type_map = None) :
         with open(file_name) as fp:
             lines = [line.rstrip('\n') for line in fp]
             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
@@ -440,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
@@ -459,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
@@ -475,7 +499,8 @@ def to_lammps_lmp(self, file_name, frame_idx = 0) :
         with open(file_name, 'w') as fp:
             fp.write(w_str)
 
-
+    @register_from_funcs.register_funcs('dump')
+    @register_from_funcs.register_funcs('lammps/dump')
     def from_lammps_dump (self,
                           file_name,
                           type_map = None,
@@ -485,13 +510,17 @@ def from_lammps_dump (self,
         self.data = dpdata.lammps.dump.system_data(lines, type_map)
         self._shift_orig_zero()
 
-
+    @register_from_funcs.register_funcs('poscar')
+    @register_from_funcs.register_funcs('contcar')
+    @register_from_funcs.register_funcs('vasp/poscar')
+    @register_from_funcs.register_funcs('vasp/contcar')
     def from_vasp_poscar(self, file_name) :
         with open(file_name) as fp:
             lines = [line.rstrip('\n') for line in fp]
             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
@@ -505,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
@@ -520,7 +550,7 @@ def to_vasp_poscar(self, file_name, frame_idx = 0) :
         with open(file_name, 'w') as fp:
             fp.write(w_str)
 
-
+    @register_from_funcs.register_funcs('qe/cp/traj')
     def from_qe_cp_traj(self,
                         prefix,
                         begin = 0,
@@ -532,15 +562,18 @@ def from_qe_cp_traj(self,
             )
         self.rot_lower_triangular()
 
-
+    @register_from_funcs.register_funcs('deepmd/npy')
     def from_deepmd_comp(self, folder, type_map = None) :
         self.data = dpdata.deepmd.comp.to_system_data(folder, type_map = type_map, labels = False)
 
+    @register_from_funcs.register_funcs('deepmd')
+    @register_from_funcs.register_funcs('deepmd/raw')
     def from_deepmd_raw(self, folder, type_map = None) :
         tmp_data = dpdata.deepmd.raw.to_system_data(folder, type_map = type_map, labels = False)
         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`.
@@ -563,12 +596,14 @@ 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`
         """
         dpdata.deepmd.raw.dump(folder, self.data) 
 
+    @register_from_funcs.register_funcs('siesta/output')
     def from_siesta_output(self, fname):
         self.data['atom_names'], \
         self.data['atom_numbs'], \
@@ -579,6 +614,7 @@ def from_siesta_output(self, fname):
             = dpdata.siesta.output.obtain_frame(fname)
         # self.rot_lower_triangular()
 
+    @register_from_funcs.register_funcs('siesta/aimd_output')
     def from_siesta_aiMD_output(self, fname):
         self.data['atom_names'], \
         self.data['atom_numbs'], \
@@ -836,38 +872,11 @@ def __init__ (self,
            return
         if file_name is None :
             return
-        if fmt == 'auto':
-            fmt = os.path.basename(file_name).split('.')[-1]
-        if fmt == 'xml' or fmt == 'XML' or fmt == 'vasp/xml' :
-            self.from_vasp_xml(file_name, begin = begin, step = step)
-        elif fmt == 'outcar' or fmt == 'OUTCAR' or fmt == 'vasp/outcar' :
-            self.from_vasp_outcar(file_name, begin = begin, step = step)
-        elif fmt == 'deepmd' or fmt == 'deepmd/raw':
-            self.from_deepmd_raw(file_name, type_map = type_map)
-        elif fmt == 'deepmd/npy':
-            self.from_deepmd_comp(file_name, type_map = type_map)
-        elif fmt == 'qe/cp/traj':
-            self.from_qe_cp_traj(file_name, begin = begin, step = step)
-        elif fmt == 'qe/pw/scf':
-            self.from_qe_pw_scf(file_name)
-        elif fmt.lower() == 'siesta/output':
-            self.from_siesta_output(file_name)
-        elif fmt.lower() == 'siesta/aimd_output':
-            self.from_siesta_aiMD_output(file_name)
-        elif fmt == 'gaussian/log':
-            self.from_gaussian_log(file_name)
-        elif fmt == 'gaussian/md':
-            self.from_gaussian_log(file_name, md=True)
-        elif fmt == 'cp2k/output':
-            self.from_cp2k_output(file_name)
-        elif fmt == 'cp2k/aimd_output':
-            self.from_cp2k_aimd_output(file_dir=file_name)
-        else :
-            raise RuntimeError('unknow data format ' + fmt)
-
+        self.from_fmt(file_name, fmt, type_map=type_map, begin= begin, step=step)
         if type_map is not None:
             self.apply_type_map(type_map)
 
+    register_from_funcs = Register()
 
     def __repr__(self):
         return self.__str__()
@@ -905,14 +914,16 @@ def has_virial(self) :
         # return ('virials' in self.data) and (len(self.data['virials']) > 0)
         return ('virials' in self.data)
 
-
+    @register_from_funcs.register_funcs('cp2k/aimd_output')
     def from_cp2k_aimd_output(self, file_dir):
         xyz_file=glob.glob("{}/*pos*.xyz".format(file_dir))[0]
         log_file=glob.glob("{}/*.log".format(file_dir))[0]
         for info_dict in Cp2kSystems(log_file, xyz_file):
             l = LabeledSystem(data=info_dict)
             self.append(l)
 
+    @register_from_funcs.register_funcs('xml')
+    @register_from_funcs.register_funcs('vasp/xml')
     def from_vasp_xml(self, file_name, begin = 0, step = 1) :
         self.data['atom_names'], \
             self.data['atom_types'], \
@@ -937,7 +948,8 @@ def from_vasp_xml(self, file_name, begin = 0, step = 1) :
         # rotate the system to lammps convention
         self.rot_lower_triangular()
 
-
+    @register_from_funcs.register_funcs('outcar')
+    @register_from_funcs.register_funcs('vasp/outcar')
     def from_vasp_outcar(self, file_name, begin = 0, step = 1) :
         # with open(file_name) as fp:
         #     lines = [line.rstrip('\n') for line in fp]
@@ -979,17 +991,18 @@ def rot_frame_lower_triangular(self, f_idx = 0) :
         self.affine_map_fv(trans, f_idx = f_idx)
         return trans
 
-
+    @register_from_funcs.register_funcs('deepmd/npy')
     def from_deepmd_comp(self, folder, type_map = None) :
         self.data = dpdata.deepmd.comp.to_system_data(folder, type_map = type_map, labels = True)
 
-
+    @register_from_funcs.register_funcs('deepmd')
+    @register_from_funcs.register_funcs('deepmd/raw')
     def from_deepmd_raw(self, folder, type_map = None) :
         tmp_data = dpdata.deepmd.raw.to_system_data(folder, type_map = type_map, labels = True)
         if tmp_data is not None :
             self.data = tmp_data
 
-
+    @register_from_funcs.register_funcs('qe/cp/traj')
     def from_qe_cp_traj(self, prefix, begin = 0, step = 1) :
         self.data, cs = dpdata.qe.traj.to_system_data(prefix + '.in', prefix, begin = begin, step = step)
         self.data['coords'] \
@@ -1001,6 +1014,7 @@ def from_qe_cp_traj(self, prefix, begin = 0, step = 1) :
         assert(cs == es), "the step key between files are not consistent" 
         self.rot_lower_triangular()
 
+    @register_from_funcs.register_funcs('qe/pw/scf')
     def from_qe_pw_scf(self, file_name) :
         self.data['atom_names'], \
             self.data['atom_numbs'], \
@@ -1013,6 +1027,7 @@ def from_qe_pw_scf(self, file_name) :
             = dpdata.qe.scf.get_frame(file_name)
         self.rot_lower_triangular()
 
+    @register_from_funcs.register_funcs('siesta/output')
     def from_siesta_output(self, file_name) :
         self.data['atom_names'], \
         self.data['atom_numbs'], \
@@ -1025,6 +1040,7 @@ def from_siesta_output(self, file_name) :
             = dpdata.siesta.output.obtain_frame(file_name)
         # self.rot_lower_triangular()
 
+    @register_from_funcs.register_funcs('siesta/aimd_output')
     def from_siesta_aiMD_output(self, file_name):
         self.data['atom_names'], \
         self.data['atom_numbs'], \
@@ -1036,14 +1052,19 @@ def from_siesta_aiMD_output(self, file_name):
         self.data['virials'] \
             = dpdata.siesta.aiMD_output.get_aiMD_frame(file_name)
 
+    @register_from_funcs.register_funcs('gaussian/log')
     def from_gaussian_log(self, file_name, md=False):
         try:
             self.data = dpdata.gaussian.log.to_system_data(file_name, md=md)
         except AssertionError:
             self.data['energies'], self.data['forces']= [], []
             self.data['nopbc'] = True
+    
+    @register_from_funcs.register_funcs('gaussian/md')
+    def from_gaussian_md(self, file_name):
+        self.from_gaussian_log(file_name, md=True)
 
-
+    @register_from_funcs.register_funcs('cp2k/output')
     def from_cp2k_output(self, file_name) :
         self.data['atom_names'], \
             self.data['atom_numbs'], \

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()