Skip to content

add experimental function register #77

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
amcadmus merged 3 commits into from
Jan 14, 2020
Merged

add experimental function register #77

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1766182
add function register
njzjz Jan 13, 2020
95c70d8
register to funcs
njzjz Jan 13, 2020
843e5ee
fix typo
njzjz Jan 13, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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`.

Expand Down
143 changes: 82 additions & 61 deletions dpdata/system.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
#%%
import os
import glob
import inspect
import numpy as np
import dpdata.lammps.lmp
import dpdata.lammps.dump
Expand All @@ -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
Expand Down Expand Up @@ -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__()
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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,
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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,
Expand All @@ -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`.
Expand All @@ -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'], \
Expand All @@ -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'], \
Expand Down Expand Up @@ -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__()
Expand Down Expand Up @@ -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'], \
Expand All @@ -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]
Expand Down Expand Up @@ -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'] \
Expand All @@ -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'], \
Expand All @@ -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'], \
Expand All @@ -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'], \
Expand All @@ -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'], \
Expand Down
9 changes: 9 additions & 0 deletions tests/test_lammps_lmp_dump.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down