From 17661825918814c414095fd1d23efbf8902fc7b8 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 13 Jan 2020 01:58:54 -0500 Subject: [PATCH 1/3] add function register --- dpdata/system.py | 122 ++++++++++++++++++++++++----------------------- 1 file changed, 62 insertions(+), 60 deletions(-) diff --git a/dpdata/system.py b/dpdata/system.py index 537ca5266..0db93a27a 100644 --- a/dpdata/system.py +++ b/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,25 @@ 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() + + def from_fmt(self, file_name, fmt, **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 __repr__(self): return self.__str__() @@ -414,7 +422,8 @@ 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] @@ -475,7 +484,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,7 +495,10 @@ 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] @@ -520,7 +533,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,10 +545,12 @@ 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 : @@ -569,6 +584,7 @@ def to_deepmd_raw(self, 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 +595,7 @@ def from_siesta_output(self, fname): = dpdata.siesta.output.obtain_frame(fname) # self.rot_lower_triangular() + @register_from_funcs.register_funcs('aimd/output') def from_siesta_aiMD_output(self, fname): self.data['atom_names'], \ self.data['atom_numbs'], \ @@ -836,38 +853,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,7 +895,7 @@ 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] @@ -913,6 +903,8 @@ def from_cp2k_aimd_output(self, file_dir): 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 +929,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 +972,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 +995,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 +1008,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 +1021,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 +1033,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'], \ From 95c70d804894722aea23e41416e4e071d8827b75 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 13 Jan 2020 13:32:15 -0500 Subject: [PATCH 2/3] register to funcs --- README.md | 10 +++++----- dpdata/system.py | 23 +++++++++++++++++++++-- tests/test_lammps_lmp_dump.py | 9 +++++++++ 3 files changed, 35 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index ae0418030..63da9a0cc 100644 --- a/README.md +++ b/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`. diff --git a/dpdata/system.py b/dpdata/system.py index 0db93a27a..b59b90774 100644 --- a/dpdata/system.py +++ b/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` diff --git a/tests/test_lammps_lmp_dump.py b/tests/test_lammps_lmp_dump.py index d20f29849..f33546485 100644 --- a/tests/test_lammps_lmp_dump.py +++ b/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() From 843e5ee14960d2267e542a4ced3601cc510f61ce Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 13 Jan 2020 13:37:32 -0500 Subject: [PATCH 3/3] fix typo --- dpdata/system.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dpdata/system.py b/dpdata/system.py index b59b90774..12cb4692e 100644 --- a/dpdata/system.py +++ b/dpdata/system.py @@ -614,7 +614,7 @@ def from_siesta_output(self, fname): = dpdata.siesta.output.obtain_frame(fname) # self.rot_lower_triangular() - @register_from_funcs.register_funcs('aimd/output') + @register_from_funcs.register_funcs('siesta/aimd_output') def from_siesta_aiMD_output(self, fname): self.data['atom_names'], \ self.data['atom_numbs'], \