From a3e76d75de53f6076254de82d18605a010dc3b00 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 3 Jul 2021 01:17:02 -0400 Subject: [PATCH 1/4] refactor and add format plugin system --- dpdata/__init__.py | 3 +- dpdata/bond_order_system.py | 48 +- dpdata/cp2k/output.py | 8 +- dpdata/format.py | 116 ++++ dpdata/plugin.py | 36 ++ dpdata/plugins/__init__.py | 19 + dpdata/plugins/abacus.py | 10 + dpdata/plugins/amber.py | 39 ++ dpdata/plugins/ase.py | 53 ++ dpdata/plugins/cp2k.py | 32 + dpdata/plugins/deepmd.py | 61 ++ dpdata/plugins/fhi_aims.py | 38 ++ dpdata/plugins/gaussian.py | 23 + dpdata/plugins/gromacs.py | 47 ++ dpdata/plugins/lammps.py | 49 ++ dpdata/plugins/list.py | 24 + dpdata/plugins/pwmat.py | 63 ++ dpdata/plugins/pymatgen.py | 47 ++ dpdata/plugins/qe.py | 45 ++ dpdata/plugins/rdkit.py | 45 ++ dpdata/plugins/siesta.py | 63 ++ dpdata/plugins/vasp.py | 110 ++++ dpdata/plugins/xyz.py | 13 + dpdata/system.py | 791 +++-------------------- plugin_example/README.md | 22 + plugin_example/dpdata_random/__init__.py | 24 + plugin_example/setup.py | 12 + tests/test_bond_order_system.py | 3 + 28 files changed, 1114 insertions(+), 730 deletions(-) create mode 100644 dpdata/format.py create mode 100644 dpdata/plugin.py create mode 100644 dpdata/plugins/__init__.py create mode 100644 dpdata/plugins/abacus.py create mode 100644 dpdata/plugins/amber.py create mode 100644 dpdata/plugins/ase.py create mode 100644 dpdata/plugins/cp2k.py create mode 100644 dpdata/plugins/deepmd.py create mode 100644 dpdata/plugins/fhi_aims.py create mode 100644 dpdata/plugins/gaussian.py create mode 100644 dpdata/plugins/gromacs.py create mode 100644 dpdata/plugins/lammps.py create mode 100644 dpdata/plugins/list.py create mode 100644 dpdata/plugins/pwmat.py create mode 100644 dpdata/plugins/pymatgen.py create mode 100644 dpdata/plugins/qe.py create mode 100644 dpdata/plugins/rdkit.py create mode 100644 dpdata/plugins/siesta.py create mode 100644 dpdata/plugins/vasp.py create mode 100644 dpdata/plugins/xyz.py create mode 100644 plugin_example/README.md create mode 100644 plugin_example/dpdata_random/__init__.py create mode 100644 plugin_example/setup.py diff --git a/dpdata/__init__.py b/dpdata/__init__.py index bd8fd8c02..001777a30 100644 --- a/dpdata/__init__.py +++ b/dpdata/__init__.py @@ -12,7 +12,8 @@ # BondOrder System has dependency on rdkit try: - import rdkit + # prevent conflict with dpdata.rdkit + import rdkit as _ USE_RDKIT = True except ModuleNotFoundError: USE_RDKIT = False diff --git a/dpdata/bond_order_system.py b/dpdata/bond_order_system.py index b590f1fcf..6b0869a9f 100644 --- a/dpdata/bond_order_system.py +++ b/dpdata/bond_order_system.py @@ -1,7 +1,6 @@ #%% # Bond Order System -from dpdata.system import Register, System, LabeledSystem, check_System -import rdkit.Chem +from dpdata.system import System, LabeledSystem, check_System, load_format import dpdata.rdkit.utils from dpdata.rdkit.sanitize import Sanitizer, SanitizeError from copy import deepcopy @@ -87,8 +86,16 @@ def __init__(self, if type_map: self.apply_type_map(type_map) - register_from_funcs = Register() - register_to_funcs = System.register_to_funcs + Register() + def from_fmt_obj(self, fmtobj, file_name, **kwargs): + mol = fmtobj.from_bond_order_system(file_name, **kwargs) + self.from_rdkit_mol(mol) + if hasattr(fmtobj.from_bond_order_system, 'post_func'): + for post_f in fmtobj.from_bond_order_system.post_func: + self.post_funcs.get_plugin(post_f)(self) + return self + + def to_fmt_obj(self, fmtobj, *args, **kwargs): + return fmtobj.to_bond_order_system(self.data, self.rdkit_mol, *args, **kwargs) def __repr__(self): return self.__str__() @@ -164,36 +171,3 @@ def from_rdkit_mol(self, rdkit_mol): self.data = dpdata.rdkit.utils.mol_to_system_data(rdkit_mol) self.data['bond_dict'] = dict([(f'{int(bond[0])}-{int(bond[1])}', bond[2]) for bond in self.data['bonds']]) self.rdkit_mol = rdkit_mol - - @register_from_funcs.register_funcs('mol') - def from_mol_file(self, file_name): - mol = rdkit.Chem.MolFromMolFile(file_name, sanitize=False, removeHs=False) - self.from_rdkit_mol(mol) - - @register_to_funcs.register_funcs("mol") - def to_mol_file(self, file_name, frame_idx=0): - assert (frame_idx < self.get_nframes()) - rdkit.Chem.MolToMolFile(self.rdkit_mol, file_name, confId=frame_idx) - - @register_from_funcs.register_funcs("sdf") - def from_sdf_file(self, file_name): - ''' - Note that it requires all molecules in .sdf file must be of the same topology - ''' - mols = [m for m in rdkit.Chem.SDMolSupplier(file_name, sanitize=False, removeHs=False)] - if len(mols) > 1: - mol = dpdata.rdkit.utils.combine_molecules(mols) - else: - mol = mols[0] - self.from_rdkit_mol(mol) - - @register_to_funcs.register_funcs("sdf") - def to_sdf_file(self, file_name, frame_idx=-1): - sdf_writer = rdkit.Chem.SDWriter(file_name) - if frame_idx == -1: - for ii in self.get_nframes(): - sdf_writer.write(self.rdkit_mol, confId=ii) - else: - assert (frame_idx < self.get_nframes()) - sdf_writer.write(self.rdkit_mol, confId=frame_idx) - sdf_writer.close() diff --git a/dpdata/cp2k/output.py b/dpdata/cp2k/output.py index ef6f1faf2..a38413719 100644 --- a/dpdata/cp2k/output.py +++ b/dpdata/cp2k/output.py @@ -278,20 +278,20 @@ def get_frames (fname) : #conver to float array and add extra dimension for nframes cell = np.array(cell) - cell = cell.astype(np.float) + cell = cell.astype(float) cell = cell[np.newaxis, :, :] coord = np.array(coord) - coord = coord.astype(np.float) + coord = coord.astype(float) coord = coord[np.newaxis, :, :] atom_symbol_list = np.array(atom_symbol_list) force = np.array(force) - force = force.astype(np.float) + force = force.astype(float) force = force[np.newaxis, :, :] # virial is not necessary if stress: stress = np.array(stress) - stress = stress.astype(np.float) + stress = stress.astype(float) stress = stress[np.newaxis, :, :] # stress to virial conversion, default unit in cp2k is GPa # note the stress is virial = stress * volume diff --git a/dpdata/format.py b/dpdata/format.py new file mode 100644 index 000000000..ee495d60f --- /dev/null +++ b/dpdata/format.py @@ -0,0 +1,116 @@ +"""Implement the format plugin system.""" +import os +from collections import abc +from abc import ABC + +from .plugin import Plugin + + +class Format(ABC): + __FormatPlugin = Plugin() + __FromPlugin = Plugin() + __ToPlugin = Plugin() + + @staticmethod + def register(key): + return Format.__FormatPlugin.register(key) + + @staticmethod + def register_from(key): + return Format.__FromPlugin.register(key) + + @staticmethod + def register_to(key): + return Format.__ToPlugin.register(key) + + @staticmethod + def get_formats(): + return Format.__FormatPlugin.plugins + + @staticmethod + def get_from_methods(): + return Format.__FromPlugin.plugins + + @staticmethod + def get_to_methods(): + return Format.__ToPlugin.plugins + + @staticmethod + def post(func_name): + def decorator(object): + if not isinstance(func_name, (list, tuple, set)): + object.post_func = (func_name,) + else: + object.post_func = func_name + return object + return decorator + + def from_system(self, file_name, **kwargs): + """System.from + + Parameters + ---------- + file_name: str + file name + + Returns + ------- + data: dict + system data + """ + raise NotImplementedError("%s doesn't support System.from" %(self.__class__.__name__)) + + def to_system(self, data, *args, **kwargs): + """System.to + + Parameters + ---------- + data: dict + system data + """ + raise NotImplementedError("%s doesn't support System.to" %(self.__class__.__name__)) + + def from_labeled_system(self, file_name, **kwargs): + raise NotImplementedError("%s doesn't support LabeledSystem.from" %(self.__class__.__name__)) + + def to_labeled_system(self, data, *args, **kwargs): + return self.to_system(data, *args, **kwargs) + + def from_bond_order_system(self, file_name, **kwargs): + raise NotImplementedError("%s doesn't support BondOrderSystem.from" %(self.__class__.__name__)) + + def to_bond_order_system(self, data, rdkit_mol, *args, **kwargs): + return self.to_system(data, *args, **kwargs) + + class MultiModes: + """File mode for MultiSystems + 0 (default): not implemented + 1: every directory under the top-level directory is a system + """ + NotImplemented = 0 + Directory = 1 + + MultiMode = MultiModes.NotImplemented + + def from_multi_systems(self, directory, **kwargs): + """MultiSystems.from + + Parameters + ---------- + directory: str + directory of system + + Returns + ------- + filenames: list[str] + 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))] + 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/plugin.py b/dpdata/plugin.py new file mode 100644 index 000000000..d000f558b --- /dev/null +++ b/dpdata/plugin.py @@ -0,0 +1,36 @@ +"""Base of plugin systems.""" + + +class Plugin: + """A class to register plugins. + + Examples + -------- + >>> Plugin = Register() + >>> @Plugin.register("xx") + def xxx(): + pass + >>> print(Plugin.plugins['xx']) + """ + def __init__(self): + self.plugins = {} + + def register(self, key): + """Register a plugin. + + Parameter + --------- + key: str + Key of the plugin. + """ + def decorator(object): + self.plugins[key] = object + return object + return decorator + + def get_plugin(self, key): + return self.plugins[key] + + def __add__(self, other): + self.plugins.update(other.plugins) + return self diff --git a/dpdata/plugins/__init__.py b/dpdata/plugins/__init__.py new file mode 100644 index 000000000..4887501ca --- /dev/null +++ b/dpdata/plugins/__init__.py @@ -0,0 +1,19 @@ +import importlib +from pathlib import Path +try: + from importlib import metadata +except ImportError: # for Python<3.8 + import importlib_metadata as metadata + +PACKAGE_BASE = "dpdata.plugins" +NOT_LOADABLE = ("__init__.py",) + +for module_file in Path(__file__).parent.glob("*.py"): + if module_file.name not in NOT_LOADABLE: + module_name = f".{module_file.stem}" + importlib.import_module(module_name, PACKAGE_BASE) + +# https://setuptools.readthedocs.io/en/latest/userguide/entry_point.html +eps = metadata.entry_points().get('dpdata.plugins', []) +for ep in eps: + plugin = ep.load() diff --git a/dpdata/plugins/abacus.py b/dpdata/plugins/abacus.py new file mode 100644 index 000000000..e9998ce6a --- /dev/null +++ b/dpdata/plugins/abacus.py @@ -0,0 +1,10 @@ +import dpdata.abacus.scf +from dpdata.format import Format + + +@Format.register("abacus/scf") +@Format.register_from("from_abacus_pw_scf") +class AbacusSCFFormat(Format): + @Format.post("rot_lower_triangular") + def from_labeled_system(self, file_name, **kwargs): + return dpdata.abacus.scf.get_frame(file_name) diff --git a/dpdata/plugins/amber.py b/dpdata/plugins/amber.py new file mode 100644 index 000000000..c7139a505 --- /dev/null +++ b/dpdata/plugins/amber.py @@ -0,0 +1,39 @@ +import dpdata.amber.md +import dpdata.amber.sqm +from dpdata.format import Format + + +@Format.register("amber/md") +@Format.register_from("from_amber_md") +class AmberMDFormat(Format): + def from_system(self, file_name=None, parm7_file=None, nc_file=None, use_element_symbols=None): + # assume the prefix is the same if the spefic name is not given + if parm7_file is None: + parm7_file = file_name + ".parm7" + if nc_file is None: + nc_file = file_name + ".nc" + return dpdata.amber.md.read_amber_traj(parm7_file=parm7_file, nc_file=nc_file, use_element_symbols=use_element_symbols, labeled=False) + + def from_labeled_system(self, file_name=None, parm7_file=None, nc_file=None, mdfrc_file=None, mden_file=None, mdout_file=None, use_element_symbols=None, **kwargs): + # assume the prefix is the same if the spefic name is not given + if parm7_file is None: + parm7_file = file_name + ".parm7" + if nc_file is None: + nc_file = file_name + ".nc" + if mdfrc_file is None: + mdfrc_file = file_name + ".mdfrc" + if mden_file is None: + mden_file = file_name + ".mden" + if mdout_file is None: + mdout_file = file_name + ".mdout" + return dpdata.amber.md.read_amber_traj(parm7_file, nc_file, mdfrc_file, mden_file, mdout_file, use_element_symbols) + + +@Format.register("sqm/out") +@Format.register_from("from_sqm_out") +class SQMOutFormat(Format): + def from_system(self, fname, **kwargs): + ''' + Read from ambertools sqm.out + ''' + return dpdata.amber.sqm.to_system_data(fname) diff --git a/dpdata/plugins/ase.py b/dpdata/plugins/ase.py new file mode 100644 index 000000000..79ba9e0b6 --- /dev/null +++ b/dpdata/plugins/ase.py @@ -0,0 +1,53 @@ +from dpdata.format import Format + + +@Format.register("ase/structure") +@Format.register_to("to_ase_structure") +class ASEStructureFormat(Format): + def to_system(self, data, **kwargs): + ''' + convert System to ASE Atom obj + + ''' + from ase import Atoms + + structures = [] + species = [data['atom_names'][tt] for tt in data['atom_types']] + + for ii in range(data['coords'].shape[0]): + structure = Atoms( + symbols=species, positions=data['coords'][ii], pbc=not data.get('nopbc', False), cell=data['cells'][ii]) + structures.append(structure) + + return structures + + def to_labeled_system(self, data, *args, **kwargs): + '''Convert System to ASE Atoms object.''' + from ase import Atoms + from ase.calculators.singlepoint import SinglePointCalculator + + structures = [] + species = [data['atom_names'][tt] for tt in data['atom_types']] + + for ii in range(data['coords'].shape[0]): + structure = Atoms( + symbols=species, + positions=data['coords'][ii], + pbc=not data.get('nopbc', False), + cell=data['cells'][ii] + ) + + results = { + 'energy': data["energies"][ii], + 'forces': data["forces"][ii] + } + if "virials" in data: + # convert to GPa as this is ase convention + v_pref = 1 * 1e4 / 1.602176621e6 + vol = structure.get_volume() + results['stress'] = data["virials"][ii] / (v_pref * vol) + + structure.calc = SinglePointCalculator(structure, **results) + structures.append(structure) + + return structures diff --git a/dpdata/plugins/cp2k.py b/dpdata/plugins/cp2k.py new file mode 100644 index 000000000..4e2933f67 --- /dev/null +++ b/dpdata/plugins/cp2k.py @@ -0,0 +1,32 @@ +import dpdata.cp2k.output +import glob +from dpdata.cp2k.output import Cp2kSystems +from dpdata.format import Format + + +@Format.register("cp2k/aimd_output") +@Format.register_from("from_cp2k_aimd_output") +class CP2KAIMDOutputFormat(Format): + def from_labeled_system(self, file_name, restart=False, **kwargs): + xyz_file = sorted(glob.glob("{}/*pos*.xyz".format(file_name)))[0] + log_file = sorted(glob.glob("{}/*.log".format(file_name)))[0] + return tuple(Cp2kSystems(log_file, xyz_file, restart)) + + +@Format.register("cp2k/output") +@Format.register_from("from_cp2k_output") +class CP2KOutputFormat(Format): + def from_labeled_system(self, file_name, restart=False, **kwargs): + data = {} + data['atom_names'], \ + data['atom_numbs'], \ + data['atom_types'], \ + data['cells'], \ + data['coords'], \ + data['energies'], \ + data['forces'], \ + tmp_virial \ + = dpdata.cp2k.output.get_frames(file_name) + if tmp_virial is not None: + data['virials'] = tmp_virial + return data diff --git a/dpdata/plugins/deepmd.py b/dpdata/plugins/deepmd.py new file mode 100644 index 000000000..0f7e8f30f --- /dev/null +++ b/dpdata/plugins/deepmd.py @@ -0,0 +1,61 @@ +import dpdata.deepmd.raw +import dpdata.deepmd.comp +import numpy as np +from dpdata.format import Format + + +@Format.register("deepmd") +@Format.register("deepmd/raw") +@Format.register_from("from_deepmd_raw") +@Format.register_to("to_deepmd_raw") +class DeePMDRawFormat(Format): + def from_system(self, file_name, type_map=None, **kwargs): + return dpdata.deepmd.raw.to_system_data(file_name, type_map=type_map, labels=False) + + def to_system(self, data, file_name, **kwargs): + """Dump the system in deepmd raw format to directory `file_name` + """ + dpdata.deepmd.raw.dump(file_name, data) + + def from_labeled_system(self, file_name, type_map, **kwargs): + return dpdata.deepmd.raw.to_system_data(file_name, type_map=type_map, labels=True) + + MultiMode = Format.MultiModes.Directory + + +@Format.register("deepmd/npy") +@Format.register_from("from_deepmd_comp") +@Format.register_from("from_deepmd_npy") +@Format.register_to("to_deepmd_comp") +@Format.register_to("to_deepmd_npy") +class DeePMDCompFormat(Format): + def from_system(self, file_name, type_map=None, **kwargs): + return dpdata.deepmd.comp.to_system_data(file_name, type_map=type_map, labels=False) + + def to_system(self, data, file_name, set_size=5000, prec=np.float32, **kwargs): + """ + Dump the system in deepmd compressed format (numpy binary) to `folder`. + + The frames are firstly split to sets, then dumped to seperated subfolders named as `folder/set.000`, `folder/set.001`, .... + + Each set contains `set_size` frames. + The last set may have less frames than `set_size`. + + Parameters + ---------- + data: dict + System data + file_name : str + The output folder + set_size : int + The size of each set. + prec: {numpy.float32, numpy.float64} + The floating point precision of the compressed data + """ + dpdata.deepmd.comp.dump( + file_name, data, set_size=set_size, comp_prec=prec) + + def from_labeled_system(self, file_name, type_map, **kwargs): + return dpdata.deepmd.comp.to_system_data(file_name, type_map=type_map, labels=True) + + MultiMode = Format.MultiModes.Directory diff --git a/dpdata/plugins/fhi_aims.py b/dpdata/plugins/fhi_aims.py new file mode 100644 index 000000000..e4b1b6239 --- /dev/null +++ b/dpdata/plugins/fhi_aims.py @@ -0,0 +1,38 @@ +import dpdata.fhi_aims.output +from dpdata.format import Format + +@Format.register("fhi_aims/md") +@Format.register_from("from_fhi_aims_output") +class FhiMDFormat(Format): + def from_labeled_system(self, file_name, md=True, begin = 0, step = 1, **kwargs): + data = {} + data['atom_names'], \ + data['atom_numbs'], \ + data['atom_types'], \ + data['cells'], \ + data['coords'], \ + data['energies'], \ + data['forces'], \ + tmp_virial, \ + = dpdata.fhi_aims.output.get_frames(file_name, md = md, begin = begin, step = step) + if tmp_virial is not None : + data['virials'] = tmp_virial + return data + +@Format.register("fhi_aims/scf") +@Format.register_from("from_fhi_aims_output") +class FhiSCFFormat(Format): + def from_labeled_system(self, file_name, **kwargs): + data = {} + data['atom_names'], \ + data['atom_numbs'], \ + data['atom_types'], \ + data['cells'], \ + data['coords'], \ + data['energies'], \ + data['forces'], \ + tmp_virial, \ + = dpdata.fhi_aims.output.get_frames(file_name, md = False, begin = 0, step = 1) + if tmp_virial is not None : + data['virials'] = tmp_virial + return data diff --git a/dpdata/plugins/gaussian.py b/dpdata/plugins/gaussian.py new file mode 100644 index 000000000..95cd3d755 --- /dev/null +++ b/dpdata/plugins/gaussian.py @@ -0,0 +1,23 @@ +import dpdata.gaussian.log +from dpdata.format import Format + + +@Format.register("gaussian/log") +@Format.register_from("from_gaussian_log") +class GaussianLogFormat(Format): + def from_labeled_system(self, file_name, md=False, **kwargs): + try: + return dpdata.gaussian.log.to_system_data(file_name, md=md) + except AssertionError: + return { + 'energies': [], + 'forces': [], + 'nopbc': True + } + + +@Format.register("gaussian/md") +@Format.register_from("from_gaussian_md") +class GaussianMDFormat(Format): + def from_labeled_system(self, file_name, **kwargs): + return GaussianLogFormat().from_labeled_system(file_name, md=True) diff --git a/dpdata/plugins/gromacs.py b/dpdata/plugins/gromacs.py new file mode 100644 index 000000000..6335c6d90 --- /dev/null +++ b/dpdata/plugins/gromacs.py @@ -0,0 +1,47 @@ +import dpdata.gromacs.gro +from dpdata.format import Format + + +@Format.register("gro") +@Format.register("gromacs/gro") +@Format.register_from("from_gromacs_gro") +@Format.register_to("to_gromacs_gro") +class PwmatOutputFormat(Format): + def from_system(self, file_name, format_atom_name=True, **kwargs): + """ + Load gromacs .gro file + + Parameters + ---------- + file_name : str + The input file name + """ + return dpdata.gromacs.gro.file_to_system_data(file_name, format_atom_name=format_atom_name) + + def to_system(self, data, file_name=None, frame_idx=-1, **kwargs): + """ + Dump the system in gromacs .gro format + + Parameters + ---------- + file_name : str or None + The output file name. If None, return the file content as a string + frame_idx : int + The index of the frame to dump + """ + assert(frame_idx < len(data['coords'])) + if frame_idx == -1: + strs = [] + for idx in range(data['coords'].shape[0]): + gro_str = dpdata.gromacs.gro.from_system_data(data, f_idx=idx) + strs.append(gro_str) + gro_str = "\n".join(strs) + else: + gro_str = dpdata.gromacs.gro.from_system_data( + data, f_idx=frame_idx) + + if file_name is None: + return gro_str + else: + with open(file_name, 'w+') as fp: + fp.write(gro_str) diff --git a/dpdata/plugins/lammps.py b/dpdata/plugins/lammps.py new file mode 100644 index 000000000..eb5dd3c8d --- /dev/null +++ b/dpdata/plugins/lammps.py @@ -0,0 +1,49 @@ +import dpdata.lammps.lmp +import dpdata.lammps.dump +from dpdata.format import Format + + +@Format.register("lmp") +@Format.register("lammps/lmp") +@Format.register_from("from_lammps_lmp") +@Format.register_to("to_lammps_lmp") +class LAMMPSLmpFormat(Format): + @Format.post("shift_orig_zero") + def from_system(self, file_name, type_map=None, **kwargs): + with open(file_name) as fp: + lines = [line.rstrip('\n') for line in fp] + return dpdata.lammps.lmp.to_system_data(lines, type_map) + + def to_system(self, data, file_name, frame_idx=0, **kwargs): + """ + Dump the system in lammps data format + + Parameters + ---------- + data: dict + System data + file_name : str + The output file name + frame_idx : int + The index of the frame to dump + """ + assert(frame_idx < len(data['coords'])) + w_str = dpdata.lammps.lmp.from_system_data(data, frame_idx) + with open(file_name, 'w') as fp: + fp.write(w_str) + + +@Format.register("dump") +@Format.register("lammps/dump") +@Format.register_from("from_lammps_dump") +@Format.register_to("to_lammps_dump") +class LAMMPSDumpFormat(Format): + @Format.post("shift_orig_zero") + def from_system(self, + file_name, + type_map=None, + begin=0, + step=1, + **kwargs): + lines = dpdata.lammps.dump.load_file(file_name, begin=begin, step=step) + return dpdata.lammps.dump.system_data(lines, type_map) diff --git a/dpdata/plugins/list.py b/dpdata/plugins/list.py new file mode 100644 index 000000000..ffe0c5381 --- /dev/null +++ b/dpdata/plugins/list.py @@ -0,0 +1,24 @@ +from dpdata.format import Format + + +@Format.register("list") +@Format.register_to("to_list") +class ListFormat(Format): + def to_system(self, data, **kwargs): + """ + convert system to list, usefull for data collection + """ + from dpdata import System, LabeledSystem + if 'forces' in data: + system = LabeledSystem(data=data) + else: + system = System(data=data) + if len(system) == 0: + return [] + if len(system) == 1: + return [system] + else: + systems = [] + for ii in range(len(system)): + systems.append(system.sub_system([ii])) + return systems diff --git a/dpdata/plugins/pwmat.py b/dpdata/plugins/pwmat.py new file mode 100644 index 000000000..0d18e466e --- /dev/null +++ b/dpdata/plugins/pwmat.py @@ -0,0 +1,63 @@ +import dpdata.pwmat.movement +import dpdata.pwmat.atomconfig +import numpy as np +from dpdata.format import Format + + +@Format.register("movement") +@Format.register("mlmd") +@Format.register("pwmat/movement") +@Format.register("pwmat/mlmd") +@Format.register_from("from_pwmat_output") +class PwmatOutputFormat(Format): + @Format.post("rot_lower_triangular") + def from_labeled_system(self, file_name, begin=0, step=1, **kwargs): + data = {} + data['atom_names'], \ + data['atom_numbs'], \ + data['atom_types'], \ + data['cells'], \ + data['coords'], \ + data['energies'], \ + data['forces'], \ + tmp_virial \ + = dpdata.pwmat.movement.get_frames(file_name, begin=begin, step=step) + if tmp_virial is not None: + data['virials'] = tmp_virial + # scale virial to the unit of eV + if 'virials' in data: + v_pref = 1 * 1e3 / 1.602176621e6 + for ii in range(data['coords'].shape[0]): + vol = np.linalg.det(np.reshape(data['cells'][ii], [3, 3])) + data['virials'][ii] *= v_pref * vol + return data + + +@Format.register("atom.config") +@Format.register("final.config") +@Format.register("pwmat/atom.config") +@Format.register("pwmat/final.config") +@Format.register_from("from_pwmat_atomconfig") +@Format.register_to("to_pwmat_atomconfig") +class PwmatAtomconfigFormat(Format): + @Format.post("rot_lower_triangular") + def from_system(self, file_name, **kwargs): + with open(file_name) as fp: + lines = [line.rstrip('\n') for line in fp] + return dpdata.pwmat.atomconfig.to_system_data(lines) + + def to_system(self, data, file_name, frame_idx=0, *args, **kwargs): + """ + Dump the system in pwmat atom.config format + + Parameters + ---------- + file_name : str + The output file name + frame_idx : int + The index of the frame to dump + """ + assert(frame_idx < len(data['coords'])) + w_str = dpdata.pwmat.atomconfig.from_system_data(data, frame_idx) + with open(file_name, 'w') as fp: + fp.write(w_str) diff --git a/dpdata/plugins/pymatgen.py b/dpdata/plugins/pymatgen.py new file mode 100644 index 000000000..8e8b63a9a --- /dev/null +++ b/dpdata/plugins/pymatgen.py @@ -0,0 +1,47 @@ +from dpdata.format import Format + + +@Format.register("pymatgen/structure") +@Format.register_to("to_pymatgen_structure") +class PyMatgenStructureFormat(Format): + def to_system(self, data, **kwargs): + """convert System to Pymatgen Structure obj + """ + structures = [] + try: + from pymatgen import Structure + except ModuleNotFoundError as e: + raise ImportError('No module pymatgen.Structure') from e + + species = [] + for name, numb in zip(data['atom_names'], data['atom_numbs']): + species.extend([name]*numb) + for ii in range(data['coords'].shape[0]): + structure = Structure( + data['cells'][ii], species, data['coords'][ii], coords_are_cartesian=True) + structures.append(structure) + return structures + + +@Format.register("pymatgen/computedstructureentry") +@Format.register_to("to_pymatgen_ComputedStructureEntry") +class PyMatgenCSEFormat(Format): + def to_labeled_system(self, data, *args, **kwargs): + """convert System to Pymagen ComputedStructureEntry obj + """ + try: + from pymatgen.entries.computed_entries import ComputedStructureEntry + except ModuleNotFoundError as e: + raise ImportError( + 'No module ComputedStructureEntry in pymatgen.entries.computed_entries') from e + + entries = [] + + for ii, structure in enumerate(PyMatgenStructureFormat().to_system(data)): + energy = data['energies'][ii] + csedata = {'forces': data['forces'][ii], + 'virials': data['virials'][ii]} + + entry = ComputedStructureEntry(structure, energy, data=csedata) + entries.append(entry) + return entries diff --git a/dpdata/plugins/qe.py b/dpdata/plugins/qe.py new file mode 100644 index 000000000..17641ec8d --- /dev/null +++ b/dpdata/plugins/qe.py @@ -0,0 +1,45 @@ +import dpdata.qe.traj +import dpdata.qe.scf +import dpdata.md.pbc +from dpdata.format import Format + +@Format.register("qe/cp/traj") +@Format.register_from("from_qe_cp_traj") +class QECPTrajFormat(Format): + @Format.post("rot_lower_triangular") + def from_system(self, file_name, begin = 0, step = 1, **kwargs): + data, _ = dpdata.qe.traj.to_system_data(file_name + '.in', file_name, begin = begin, step = step) + data['coords'] \ + = dpdata.md.pbc.apply_pbc(data['coords'], + data['cells'], + ) + return data + + @Format.post("rot_lower_triangular") + def from_labeled_system(self, file_name, begin = 0, step = 1, **kwargs): + data, cs = dpdata.qe.traj.to_system_data(file_name + '.in', file_name, begin = begin, step = step) + data['coords'] \ + = dpdata.md.pbc.apply_pbc(data['coords'], + data['cells'], + ) + data['energies'], data['forces'], es \ + = dpdata.qe.traj.to_system_label(file_name + '.in', file_name, begin = begin, step = step) + assert(cs == es), "the step key between files are not consistent" + return data + +@Format.register("qe/pw/scf") +@Format.register_from("from_qe_pw_scf") +class QECPTrajFormat(Format): + @Format.post("rot_lower_triangular") + def from_labeled_system(self, file_name, **kwargs): + data = {} + data['atom_names'], \ + data['atom_numbs'], \ + data['atom_types'], \ + data['cells'], \ + data['coords'], \ + data['energies'], \ + data['forces'], \ + data['virials'], \ + = dpdata.qe.scf.get_frame(file_name) + return data diff --git a/dpdata/plugins/rdkit.py b/dpdata/plugins/rdkit.py new file mode 100644 index 000000000..bc19bf3a1 --- /dev/null +++ b/dpdata/plugins/rdkit.py @@ -0,0 +1,45 @@ +from dpdata.format import Format +try: + import rdkit.Chem +except ModuleNotFoundError: + pass +import dpdata.rdkit.utils + + +@Format.register("mol") +@Format.register_from("from_mol_file") +@Format.register_to("to_mol_file") +class MolFormat(Format): + def from_bond_order_system(self, file_name, **kwargs): + return rdkit.Chem.MolFromMolFile(file_name, sanitize=False, removeHs=False) + + + def to_bond_order_system(self, mol, data, file_name, frame_idx=0, **kwargs): + assert (frame_idx < self.get_nframes()) + rdkit.Chem.MolToMolFile(mol, file_name, confId=frame_idx) + + +@Format.register("sdf") +@Format.register_from("from_sdf_file") +@Format.register_to("to_sdf_file") +class SdfFormat(Format): + def from_bond_order_system(self, file_name, **kwargs): + ''' + Note that it requires all molecules in .sdf file must be of the same topology + ''' + mols = [m for m in rdkit.Chem.SDMolSupplier(file_name, sanitize=False, removeHs=False)] + if len(mols) > 1: + mol = dpdata.rdkit.utils.combine_molecules(mols) + else: + mol = mols[0] + return mol + + def to_bond_order_system(self, mol, data, file_name, frame_idx=-1, **kwargs): + sdf_writer = rdkit.Chem.SDWriter(file_name) + if frame_idx == -1: + for ii in range(data['coords'].shape[0]): + sdf_writer.write(mol, confId=ii) + else: + assert (frame_idx < data['coords'].shape[0]) + sdf_writer.write(mol, confId=frame_idx) + sdf_writer.close() \ No newline at end of file diff --git a/dpdata/plugins/siesta.py b/dpdata/plugins/siesta.py new file mode 100644 index 000000000..c6bb55399 --- /dev/null +++ b/dpdata/plugins/siesta.py @@ -0,0 +1,63 @@ +import dpdata.siesta.output +import dpdata.siesta.aiMD_output +from dpdata.format import Format + + +@Format.register("siesta/output") +@Format.register_from("from_siesta_output") +class SiestaOutputFormat(Format): + def from_system(self, file_name, **kwargs): + data = {} + data['atom_names'], \ + data['atom_numbs'], \ + data['atom_types'], \ + data['cells'], \ + data['coords'], \ + _e, \ + _f, \ + _v \ + = dpdata.siesta.output.obtain_frame(file_name) + return data + + def from_labeled_system(self, file_name, **kwargs): + data = {} + data['atom_names'], \ + data['atom_numbs'], \ + data['atom_types'], \ + data['cells'], \ + data['coords'], \ + data['energies'], \ + data['forces'], \ + data['virials'] \ + = dpdata.siesta.output.obtain_frame(file_name) + return data + + +@Format.register("siesta/aimd_output") +@Format.register_from("from_siesta_aiMD_output") +class SiestaOutputFormat(Format): + def from_system(self, file_name, **kwargs): + data = {} + data['atom_names'], \ + data['atom_numbs'], \ + data['atom_types'], \ + data['cells'], \ + data['coords'], \ + _e, \ + _f, \ + _v \ + = dpdata.siesta.aiMD_output.get_aiMD_frame(file_name) + return data + + def from_labeled_system(self, file_name, **kwargs): + data = {} + data['atom_names'], \ + data['atom_numbs'], \ + data['atom_types'], \ + data['cells'], \ + data['coords'], \ + data['energies'], \ + data['forces'], \ + data['virials'] \ + = dpdata.siesta.aiMD_output.get_aiMD_frame(file_name) + return data diff --git a/dpdata/plugins/vasp.py b/dpdata/plugins/vasp.py new file mode 100644 index 000000000..fca08f15d --- /dev/null +++ b/dpdata/plugins/vasp.py @@ -0,0 +1,110 @@ +import dpdata.vasp.poscar +import dpdata.vasp.xml +import dpdata.vasp.outcar +import numpy as np +from dpdata.format import Format + + +@Format.register("poscar") +@Format.register("contcar") +@Format.register("vasp/poscar") +@Format.register("vasp/contcar") +@Format.register_from("from_vasp_poscar") +@Format.register_to("to_vasp_poscar") +class VASPPoscarFormat(Format): + @Format.post("rot_lower_triangular") + def from_system(self, file_name, **kwargs): + with open(file_name) as fp: + lines = [line.rstrip('\n') for line in fp] + return dpdata.vasp.poscar.to_system_data(lines) + + def to_system(self, data, file_name, frame_idx=0, **kwargs): + """ + Dump the system in vasp POSCAR format + + Parameters + ---------- + file_name : str + The output file name + frame_idx : int + The index of the frame to dump + """ + w_str = VASPStringFormat().to_system(data, frame_idx=frame_idx) + with open(file_name, 'w') as fp: + fp.write(w_str) + + +@Format.register("vasp/string") +@Format.register_to("to_vasp_string") +class VASPStringFormat(Format): + def to_system(self, data, frame_idx=0, **kwargs): + """ + Dump the system in vasp POSCAR format string + + Parameters + ---------- + frame_idx : int + The index of the frame to dump + """ + assert(frame_idx < len(data['coords'])) + return dpdata.vasp.poscar.from_system_data(data, frame_idx) + + +# rotate the system to lammps convention +@Format.register("outcar") +@Format.register("vasp/outcar") +@Format.register_from("from_vasp_outcar") +class VASPOutcarFormat(Format): + @Format.post("rot_lower_triangular") + def from_labeled_system(self, file_name, begin=0, step=1, **kwargs): + data = {} + data['atom_names'], \ + data['atom_numbs'], \ + data['atom_types'], \ + data['cells'], \ + data['coords'], \ + data['energies'], \ + data['forces'], \ + tmp_virial, \ + = dpdata.vasp.outcar.get_frames(file_name, begin=begin, step=step) + if tmp_virial is not None: + data['virials'] = tmp_virial + # scale virial to the unit of eV + if 'virials' in data: + v_pref = 1 * 1e3 / 1.602176621e6 + for ii in range(data['cells'].shape[0]): + vol = np.linalg.det(np.reshape(data['cells'][ii], [3, 3])) + data['virials'][ii] *= v_pref * vol + return data + + +# rotate the system to lammps convention +@Format.register("xml") +@Format.register("vasp/xml") +@Format.register_from("from_vasp_xml") +class VASPXMLFormat(Format): + @Format.post("rot_lower_triangular") + def from_labeled_system(self, file_name, begin=0, step=1, **kwargs): + data = {} + data['atom_names'], \ + data['atom_types'], \ + data['cells'], \ + data['coords'], \ + data['energies'], \ + data['forces'], \ + data['virials'], \ + = dpdata.vasp.xml.analyze(file_name, type_idx_zero=True, begin=begin, step=step) + data['atom_numbs'] = [] + for ii in range(len(data['atom_names'])): + data['atom_numbs'].append(sum(data['atom_types'] == ii)) + # the vasp xml assumes the direct coordinates + # apply the transform to the cartesan coordinates + for ii in range(data['cells'].shape[0]): + data['coords'][ii] = np.matmul( + data['coords'][ii], data['cells'][ii]) + # scale virial to the unit of eV + v_pref = 1 * 1e3 / 1.602176621e6 + for ii in range(data['cells'].shape[0]): + vol = np.linalg.det(np.reshape(data['cells'][ii], [3, 3])) + data['virials'][ii] *= v_pref * vol + return data diff --git a/dpdata/plugins/xyz.py b/dpdata/plugins/xyz.py new file mode 100644 index 000000000..a3fa7b27c --- /dev/null +++ b/dpdata/plugins/xyz.py @@ -0,0 +1,13 @@ +from dpdata.xyz.quip_gap_xyz import QuipGapxyzSystems +from dpdata.format import Format + + +@Format.register("quip/gap/xyz") +@Format.register_from("from_quip_gap_xyz_file") +class DeePMDRawFormat(Format): + def from_labeled_system(self, data, **kwargs): + return data + + def from_multi_systems(self, file_name, **kwargs): + # here directory is the file_name + return QuipGapxyzSystems(file_name) \ No newline at end of file diff --git a/dpdata/system.py b/dpdata/system.py index a84f31c59..5ca0f24bd 100644 --- a/dpdata/system.py +++ b/dpdata/system.py @@ -3,49 +3,27 @@ import glob import inspect import numpy as np -import dpdata.lammps.lmp -import dpdata.lammps.dump -import dpdata.vasp.poscar -import dpdata.vasp.xml -import dpdata.vasp.outcar -import dpdata.deepmd.raw -import dpdata.deepmd.comp -import dpdata.qe.traj -import dpdata.qe.scf -import dpdata.abacus.scf -import dpdata.siesta.output -import dpdata.siesta.aiMD_output import dpdata.md.pbc -import dpdata.gaussian.log -import dpdata.amber.md -import dpdata.amber.sqm -import dpdata.cp2k.output -from dpdata.cp2k.output import Cp2kSystems -import dpdata.pwmat.movement -import dpdata.pwmat.atomconfig -import dpdata.fhi_aims.output -import dpdata.gromacs.gro from copy import deepcopy from monty.json import MSONable from monty.serialization import loadfn,dumpfn from dpdata.periodic_table import Element -from dpdata.xyz.quip_gap_xyz import QuipGapxyzSystems from dpdata.amber.mask import pick_by_amber_mask, load_param_file +# ensure all plugins are loaded! +import dpdata.plugins +from dpdata.plugin import Plugin +from dpdata.format import Format -class Register: - def __init__(self): - self.funcs = {} - - def register_funcs(self, fmt): - def decorator(func): - self.funcs[fmt] = func - return func - return decorator - - def __add__(self, other): - self.funcs.update(other.funcs) - return self +def load_format(fmt): + fmt = fmt.lower() + formats = Format.get_formats() + if fmt in formats: + return formats[fmt]() + raise NotImplementedError( + "Unsupported data format %s. Supported formats: %s" % ( + fmt, " ".join(formats) + )) class System (MSONable) : @@ -131,32 +109,32 @@ def __init__ (self, if type_map is not None: self.apply_type_map(type_map) - register_from_funcs = Register() - register_to_funcs = Register() + post_funcs = Plugin() 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) + return self.from_fmt_obj(load_format(fmt), file_name, **kwargs) + + def from_fmt_obj(self, fmtobj, file_name, **kwargs): + data = fmtobj.from_system(file_name, **kwargs) + if data: + if isinstance(data, (list, tuple)): + for dd in data: + self.append(System(data=dd)) + else: + self.data = {**self.data, **data} + if hasattr(fmtobj.from_system, 'post_func'): + for post_f in fmtobj.from_system.post_func: + self.post_funcs.get_plugin(post_f)(self) + return self 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: %s' % (fmt, " ".join(to_funcs))) + return self.to_fmt_obj(load_format(fmt), *args, **kwargs) + + def to_fmt_obj(self, fmtobj, *args, **kwargs): + return fmtobj.to_system(self.data, *args, **kwargs) def __repr__(self): return self.__str__() @@ -243,21 +221,6 @@ 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 - """ - if len(self)==0: - return [] - if len(self)==1: - return [self] - else: - systems=[] - for ii in range(len(self)): - systems.append(self.sub_system([ii])) - return systems - @staticmethod def load(filename): """rebuild System obj. from .json or .yaml file """ @@ -492,276 +455,13 @@ def remove_pbc(self, protect_layer = 9): self.data['coords'][ff] = self.data['coords'][ff] + np.tile(shift, [natoms, 1]) self.data['cells'][ff] = cell_size * np.eye(3) - - @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 - - ''' - structures=[] - try: - from pymatgen import Structure - except: - raise ImportError('No module pymatgen.Structure') - - for system in self.to_list(): - species=[] - for name,numb in zip(system.data['atom_names'],system.data['atom_numbs']): - species.extend([name]*numb) - structure=Structure(system.data['cells'][0],species,system.data['coords'][0],coords_are_cartesian=True) - structures.append(structure) - return structures - - - @register_to_funcs.register_funcs("ase/structure") - def to_ase_structure(self): - ''' - convert System to ASE Atom obj - - ''' - from ase import Atoms - - structures=[] - - for system in self.to_list(): - species=[system.data['atom_names'][tt] for tt in system.data['atom_types']] - structure=Atoms(symbols=species,positions=system.data['coords'][0],pbc=True,cell=system.data['cells'][0]) - 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 - - Parameters - ---------- - file_name : str - The output file name - frame_idx : int - The index of the frame to dump - """ - assert(frame_idx < len(self.data['coords'])) - w_str = dpdata.lammps.lmp.from_system_data(self.data, frame_idx) - 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, - begin = 0, - step = 1) : - lines = dpdata.lammps.dump.load_file(file_name, begin = begin, step = step) - 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 - - Parameters - ---------- - frame_idx : int - The index of the frame to dump - """ - assert(frame_idx < len(self.data['coords'])) - 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 - - Parameters - ---------- - file_name : str - The output file name - frame_idx : int - The index of the frame to dump - """ - w_str=self.to_vasp_string( frame_idx= frame_idx ) - 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, - step = 1) : - self.data, _ = dpdata.qe.traj.to_system_data(prefix + '.in', prefix, begin = begin, step = step) - self.data['coords'] \ - = dpdata.md.pbc.apply_pbc(self.data['coords'], - self.data['cells'], - ) - 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_from_funcs.register_funcs("gro") - @register_from_funcs.register_funcs("gromacs/gro") - def from_gromacs_gro(self, file_name, format_atom_name=True) : - """ - Load gromacs .gro file - - Parameters - ---------- - file_name : str - The input file name - """ - self.data = dpdata.gromacs.gro.file_to_system_data(file_name, format_atom_name=format_atom_name) - - @register_to_funcs.register_funcs("gromacs/gro") - def to_gromacs_gro(self, file_name=None, frame_idx=-1): - """ - Dump the system in gromacs .gro format - - Parameters - ---------- - file_name : str or None - The output file name. If None, return the file content as a string - frame_idx : int - The index of the frame to dump - """ - assert(frame_idx < len(self.data['coords'])) - if frame_idx == -1: - strs = [] - for idx in range(self.get_nframes()): - gro_str = dpdata.gromacs.gro.from_system_data(self.data, f_idx=idx) - strs.append(gro_str) - gro_str = "\n".join(strs) - else: - gro_str = dpdata.gromacs.gro.from_system_data(self.data, f_idx=frame_idx) - - if file_name is None: - return gro_str - else: - with open(file_name, 'w+') as fp: - fp.write(gro_str) - - @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`. - - The frames are firstly split to sets, then dumped to seperated subfolders named as `folder/set.000`, `folder/set.001`, .... - - Each set contains `set_size` frames. - The last set may have less frames than `set_size`. - - Parameters - ---------- - folder : str - The output folder - set_size : int - The size of each set. - prec: {numpy.float32, numpy.float64} - The floating point precision of the compressed data - """ - dpdata.deepmd.comp.dump(folder, self.data, - 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('sqm/out') - def from_sqm_out(self, fname): - ''' - Read from ambertools sqm.out - ''' - self.data = dpdata.amber.sqm.to_system_data(fname) - - @register_from_funcs.register_funcs('siesta/output') - def from_siesta_output(self, fname): - self.data['atom_names'], \ - self.data['atom_numbs'], \ - self.data['atom_types'], \ - self.data['cells'], \ - self.data['coords'], \ - _e, _f, _v \ - = 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'], \ - self.data['atom_types'], \ - self.data['cells'], \ - self.data['coords'], \ - _e, _f, _v \ - = dpdata.siesta.aiMD_output.get_aiMD_frame(fname) - @register_from_funcs.register_funcs('atom.config') - @register_from_funcs.register_funcs('final.config') - @register_from_funcs.register_funcs('pwmat/atom.config') - @register_from_funcs.register_funcs('pwmat/final.config') - def from_pwmat_atomconfig(self, file_name) : - with open(file_name) as fp: - lines = [line.rstrip('\n') for line in fp] - self.data = dpdata.pwmat.atomconfig.to_system_data(lines) - self.rot_lower_triangular() - - @register_to_funcs.register_funcs("pwmat/atom.config") - def to_pwmat_atomconfig(self, file_name, frame_idx = 0) : - """ - Dump the system in pwmat atom.config format - - Parameters - ---------- - file_name : str - The output file name - frame_idx : int - The index of the frame to dump - """ - assert(frame_idx < len(self.data['coords'])) - w_str = dpdata.pwmat.atomconfig.from_system_data(self.data, frame_idx) - with open(file_name, 'w') as fp: - fp.write(w_str) - - def affine_map(self, trans, f_idx = 0) : assert(np.linalg.det(trans) != 0) self.data['cells'][f_idx] = np.matmul(self.data['cells'][f_idx], trans) self.data['coords'][f_idx] = np.matmul(self.data['coords'][f_idx], trans) + @post_funcs.register("shift_orig_zero") def _shift_orig_zero(self) : for ff in self.data['coords'] : for ii in ff : @@ -769,7 +469,7 @@ def _shift_orig_zero(self) : self.data['orig'] = self.data['orig'] - self.data['orig'] assert((np.zeros([3]) == self.data['orig']).all()) - + @post_funcs.register("rot_lower_triangular") def rot_lower_triangular(self) : for ii in range(self.get_nframes()) : self.rot_frame_lower_triangular(ii) @@ -1067,15 +767,6 @@ def pick_by_amber_mask(self, param, maskstr, pass_coords=False, nopbc=None): idx = pick_by_amber_mask(parm, maskstr) return self.pick_atom_idx(idx, nopbc=nopbc) - @register_from_funcs.register_funcs('amber/md') - def from_amber_md(self, file_name=None, parm7_file=None, nc_file=None, use_element_symbols=None): - # assume the prefix is the same if the spefic name is not given - if parm7_file is None: - parm7_file = file_name + ".parm7" - if nc_file is None: - nc_file = file_name + ".nc" - self.data = dpdata.amber.md.read_amber_traj(parm7_file=parm7_file, nc_file=nc_file, use_element_symbols=use_element_symbols, labeled=False) - def get_cell_perturb_matrix(cell_pert_fraction): if cell_pert_fraction<0: raise RuntimeError('cell_pert_fraction can not be negative') @@ -1162,7 +853,7 @@ def __init__ (self, - ``pwmat/out.mlmd``: pwmat scf output file type_map : list of str - Needed by formats deepmd/raw and deepmd/npy. Maps atom type to name. The atom with type `ii` is mapped to `type_map[ii]`. + Maps atom type to name. The atom with type `ii` is mapped to `type_map[ii]`. If not provided the atom names are assigned to `'Type_1'`, `'Type_2'`, `'Type_3'`... begin : int The beginning frame when loading MD trajectory. @@ -1182,8 +873,23 @@ def __init__ (self, if type_map is not None: self.apply_type_map(type_map) - register_from_funcs = Register() - register_to_funcs = System.register_to_funcs + Register() + post_funcs = Plugin() + System.post_funcs + + def from_fmt_obj(self, fmtobj, file_name, **kwargs): + data = fmtobj.from_labeled_system(file_name, **kwargs) + if data: + if isinstance(data, (list, tuple)): + for dd in data: + self.append(LabeledSystem(data=dd)) + else: + self.data = {**self.data, **data} + if hasattr(fmtobj.from_labeled_system, 'post_func'): + for post_f in fmtobj.from_labeled_system.post_func: + self.post_funcs.get_plugin(post_f)(self) + return self + + def to_fmt_obj(self, fmtobj, *args, **kwargs): + return fmtobj.to_labeled_system(self.data, *args, **kwargs) def __repr__(self): return self.__str__() @@ -1221,255 +927,22 @@ 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, restart=False): - xyz_file=sorted(glob.glob("{}/*pos*.xyz".format(file_dir)))[0] - log_file=sorted(glob.glob("{}/*.log".format(file_dir)))[0] - for info_dict in Cp2kSystems(log_file, xyz_file, restart): - l = LabeledSystem(data=info_dict) - self.append(l) - - # @register_from_funcs.register_funcs('cp2k/restart_aimd_output') - # def from_cp2k_aimd_output(self, file_dir, restart=True): - # xyz_file = sorted(glob.glob("{}/*pos*.xyz".format(file_dir)))[0] - # log_file = sorted(glob.glob("{}/*.log".format(file_dir)))[0] - # for info_dict in Cp2kSystems(log_file, xyz_file, restart): - # l = LabeledSystem(data=info_dict) - # self.append(l) - - @register_from_funcs.register_funcs('fhi_aims/md') - def from_fhi_aims_output(self, file_name, md=True, begin=0, step =1): - self.data['atom_names'], \ - self.data['atom_numbs'], \ - self.data['atom_types'], \ - self.data['cells'], \ - self.data['coords'], \ - self.data['energies'], \ - self.data['forces'], \ - tmp_virial, \ - = dpdata.fhi_aims.output.get_frames(file_name, md = md, begin = begin, step = step) - if tmp_virial is not None : - self.data['virials'] = tmp_virial - - @register_from_funcs.register_funcs('fhi_aims/scf') - def from_fhi_aims_output(self, file_name ): - self.data['atom_names'], \ - self.data['atom_numbs'], \ - self.data['atom_types'], \ - self.data['cells'], \ - self.data['coords'], \ - self.data['energies'], \ - self.data['forces'], \ - tmp_virial, \ - = dpdata.fhi_aims.output.get_frames(file_name, md = False, begin = 0, step = 1) - if tmp_virial is not None : - self.data['virials'] = tmp_virial - - @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'], \ - self.data['cells'], \ - self.data['coords'], \ - self.data['energies'], \ - self.data['forces'], \ - self.data['virials'], \ - = dpdata.vasp.xml.analyze(file_name, type_idx_zero = True, begin = begin, step = step) - self.data['atom_numbs'] = [] - for ii in range(len(self.data['atom_names'])) : - self.data['atom_numbs'].append(sum(self.data['atom_types'] == ii)) - # the vasp xml assumes the direct coordinates - # apply the transform to the cartesan coordinates - for ii in range(self.get_nframes()) : - self.data['coords'][ii] = np.matmul(self.data['coords'][ii], self.data['cells'][ii]) - # scale virial to the unit of eV - v_pref = 1 * 1e3 / 1.602176621e6 - for ii in range (self.get_nframes()) : - vol = np.linalg.det(np.reshape(self.data['cells'][ii], [3,3])) - self.data['virials'][ii] *= v_pref * vol - # 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] - self.data['atom_names'], \ - self.data['atom_numbs'], \ - self.data['atom_types'], \ - self.data['cells'], \ - self.data['coords'], \ - self.data['energies'], \ - self.data['forces'], \ - tmp_virial, \ - = dpdata.vasp.outcar.get_frames(file_name, begin = begin, step = step) - if tmp_virial is not None : - self.data['virials'] = tmp_virial - # scale virial to the unit of eV - if 'virials' in self.data : - v_pref = 1 * 1e3 / 1.602176621e6 - for ii in range (self.get_nframes()) : - vol = np.linalg.det(np.reshape(self.data['cells'][ii], [3,3])) - self.data['virials'][ii] *= v_pref * vol - # rotate the system to lammps convention - self.rot_lower_triangular() - - def affine_map_fv(self, trans, f_idx) : assert(np.linalg.det(trans) != 0) self.data['forces'][f_idx] = np.matmul(self.data['forces'][f_idx], trans) if self.has_virial(): self.data['virials'][f_idx] = np.matmul(trans.T, np.matmul(self.data['virials'][f_idx], trans)) - + @post_funcs.register("rot_lower_triangular") def rot_lower_triangular(self) : for ii in range(self.get_nframes()) : self.rot_frame_lower_triangular(ii) - def rot_frame_lower_triangular(self, f_idx = 0) : trans = System.rot_frame_lower_triangular(self, f_idx = f_idx) 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'] \ - = dpdata.md.pbc.apply_pbc(self.data['coords'], - self.data['cells'], - ) - self.data['energies'], self.data['forces'], es \ - = dpdata.qe.traj.to_system_label(prefix + '.in', prefix, begin = begin, step = step) - 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'], \ - self.data['atom_types'], \ - self.data['cells'], \ - self.data['coords'], \ - self.data['energies'], \ - self.data['forces'], \ - self.data['virials'], \ - = dpdata.qe.scf.get_frame(file_name) - self.rot_lower_triangular() - - @register_from_funcs.register_funcs('abacus/scf') - def from_abacus_pw_scf(self, file_name) : - self.data = dpdata.abacus.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'], \ - self.data['atom_types'], \ - self.data['cells'], \ - self.data['coords'], \ - self.data['energies'], \ - self.data['forces'], \ - self.data['virials'] \ - = 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'], \ - self.data['atom_types'], \ - self.data['cells'], \ - self.data['coords'], \ - self.data['energies'], \ - self.data['forces'], \ - 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('amber/md') - def from_amber_md(self, file_name=None, parm7_file=None, nc_file=None, mdfrc_file=None, mden_file=None, mdout_file=None, use_element_symbols=None): - # assume the prefix is the same if the spefic name is not given - if parm7_file is None: - parm7_file = file_name + ".parm7" - if nc_file is None: - nc_file = file_name + ".nc" - if mdfrc_file is None: - mdfrc_file = file_name + ".mdfrc" - if mden_file is None: - mden_file = file_name + ".mden" - if mdout_file is None: - mdout_file = file_name + ".mdout" - self.data = dpdata.amber.md.read_amber_traj(parm7_file, nc_file, mdfrc_file, mden_file, mdout_file, use_element_symbols) - - @register_from_funcs.register_funcs('cp2k/output') - def from_cp2k_output(self, file_name) : - self.data['atom_names'], \ - self.data['atom_numbs'], \ - self.data['atom_types'], \ - self.data['cells'], \ - self.data['coords'], \ - self.data['energies'], \ - self.data['forces'], \ - tmp_virial \ - = dpdata.cp2k.output.get_frames(file_name) - if tmp_virial is not None: - self.data['virials'] = tmp_virial - @register_from_funcs.register_funcs('movement') - @register_from_funcs.register_funcs('MOVEMENT') - @register_from_funcs.register_funcs('mlmd') - @register_from_funcs.register_funcs('MLMD') - @register_from_funcs.register_funcs('pwmat/movement') - @register_from_funcs.register_funcs('pwmat/MOVEMENT') - @register_from_funcs.register_funcs('pwmat/mlmd') - @register_from_funcs.register_funcs('pwmat/MLMD') - def from_pwmat_output(self, file_name, begin = 0, step = 1) : - self.data['atom_names'], \ - self.data['atom_numbs'], \ - self.data['atom_types'], \ - self.data['cells'], \ - self.data['coords'], \ - self.data['energies'], \ - self.data['forces'], \ - tmp_virial, \ - = dpdata.pwmat.movement.get_frames(file_name, begin = begin, step = step) - if tmp_virial is not None : - self.data['virials'] = tmp_virial - # scale virial to the unit of eV - if 'virials' in self.data : - v_pref = 1 * 1e3 / 1.602176621e6 - for ii in range (self.get_nframes()) : - vol = np.linalg.det(np.reshape(self.data['cells'][ii], [3,3])) - self.data['virials'][ii] *= v_pref * vol - # rotate the system to lammps convention - self.rot_lower_triangular() - - def sub_system(self, f_idx) : """ Construct a subsystem from the system @@ -1492,38 +965,6 @@ def sub_system(self, f_idx) : tmp_sys.data['virials'] = self.data['virials'][f_idx].reshape(-1, 3, 3) return tmp_sys - @register_to_funcs.register_funcs("ase/structure") - def to_ase_structure(self): - '''Convert System to ASE Atoms object.''' - from ase import Atoms - from ase.calculators.singlepoint import SinglePointCalculator - - structures = [] - - for system in self.to_list(): - species=[system.data['atom_names'][tt] for tt in system.data['atom_types']] - structure=Atoms( - symbols=species, - positions=system.data['coords'][0], - pbc=True, - cell=system.data['cells'][0] - ) - - results = { - 'energy': system.data["energies"][0], - 'forces': system.data["forces"][0] - } - if "virials" in system.data: - # convert to GPa as this is ase convention - v_pref = 1 * 1e4 / 1.602176621e6 - vol = structure.get_volume() - results['stress'] = system.data["virials"][0] / (v_pref * vol) - - structure.calc = SinglePointCalculator(structure, **results) - structures.append(structure) - - return structures - def append(self, system): """ Append a system to this system @@ -1564,27 +1005,6 @@ def shuffle(self): self.data[ii] = self.data[ii][idx] return idx - def to_pymatgen_ComputedStructureEntry(self): - ''' - convert System to Pymagen ComputedStructureEntry obj - - ''' - try: - from pymatgen.entries.computed_entries import ComputedStructureEntry - except: - raise ImportError('No module ComputedStructureEntry in pymatgen.entries.computed_entries') - - entries=[] - for system in self.to_list(): - structure=system.to_pymatgen_structure()[0] - energy=system.data['energies'][0] - data={'forces':system.data['forces'][0], - 'virials':system.data['virials'][0]} - - entry=ComputedStructureEntry(structure,energy,data=data) - entries.append(entry) - return entries - def correction(self, hl_sys): """Get energy and force correction between self and a high-level LabeledSystem. The self's coordinates will be kept, but energy and forces will be replaced by @@ -1651,6 +1071,21 @@ def __init__(self, *systems,type_map=None): self.atom_names = [] self.append(*systems) + def from_fmt_obj(self, fmtobj, directory, labeled=True, **kwargs): + for dd in fmtobj.from_multi_systems(directory, **kwargs): + if labeled: + system = LabeledSystem().from_fmt_obj(fmtobj, dd, **kwargs) + else: + system = System().from_fmt_obj(fmtobj, dd, **kwargs) + system.sort_atom_names() + self.append(system) + return self + + def to_fmt_obj(self, fmtobj, directory, *args, **kwargs): + for fn, ss in zip(fmtobj.to_multi_systems(self.systems.keys(), directory, **kwargs), self.systems.values()): + ss.to_fmt_obj(fmtobj, fn, *args, **kwargs) + return self + def __getitem__(self, key): """Returns proerty stored in System by key or by idx""" if isinstance(key, int): @@ -1676,9 +1111,9 @@ def __add__(self, others) : raise RuntimeError("Unspported data structure") @classmethod - def from_file(cls,file_name,fmt): + def from_file(cls,file_name,fmt, **kwargs): multi_systems = cls() - multi_systems.load_systems_from_file(file_name=file_name,fmt=fmt) + multi_systems.load_systems_from_file(file_name=file_name,fmt=fmt, **kwargs) return multi_systems @classmethod @@ -1690,15 +1125,9 @@ def from_dir(cls,dir_name, file_name, fmt='auto', type_map=None): return multi_systems - def load_systems_from_file(self, file_name=None, fmt=None): - if file_name is not None: - if fmt is None: - raise RuntimeError("must specify file format for file {}".format(file_name)) - elif fmt == 'quip/gap/xyz' or 'xyz': - self.from_quip_gap_xyz_file(file_name) - else: - raise RuntimeError("unknown file format for file {} format {},now supported 'quip/gap/xyz'".format(file_name, fmt)) - + def load_systems_from_file(self, file_name=None, fmt=None, **kwargs): + fmt = fmt.lower() + return self.from_fmt_obj(load_format(fmt), file_name, **kwargs) def get_nframes(self) : """Returns number of frames in all systems""" @@ -1755,50 +1184,6 @@ def check_atom_names(self, system): system.add_atom_names(new_in_self) system.sort_atom_names(type_map=self.atom_names) - def from_quip_gap_xyz_file(self,file_name): - # quip_gap_xyz_systems = QuipGapxyzSystems(file_name) - # print(next(quip_gap_xyz_systems)) - for info_dict in QuipGapxyzSystems(file_name): - system=LabeledSystem(data=info_dict) - system.sort_atom_names() - self.append(system) - - - def to_deepmd_raw(self, folder) : - """ - Dump systems in deepmd raw format to `folder` for each system. - """ - for system_name, system in self.systems.items(): - system.to_deepmd_raw(os.path.join(folder, system_name)) - - def to_deepmd_npy(self, folder, set_size = 5000, prec=np.float32) : - """ - Dump the system in deepmd compressed format (numpy binary) to `folder` for each system. - - Parameters - ---------- - folder : str - The output folder - set_size : int - The size of each set. - prec: {numpy.float32, numpy.float64} - The floating point precision of the compressed data - """ - for system_name, system in self.systems.items(): - system.to_deepmd_npy(os.path.join(folder, system_name), - set_size = set_size, - prec = prec) - - def from_deepmd_raw(self, folder): - for dd in os.listdir(folder): - self.append(LabeledSystem(os.path.join(folder, dd), fmt='deepmd/raw')) - return self - - def from_deepmd_npy(self, folder): - for dd in os.listdir(folder): - self.append(LabeledSystem(os.path.join(folder, dd), fmt='deepmd/npy')) - return self - def predict(self, dp): import deepmd.DeepPot as DeepPot if not isinstance(dp, DeepPot): @@ -1829,6 +1214,36 @@ def pick_atom_idx(self, idx, nopbc=None): return new_sys +def add_format_methods(): + """Add format methods to System, LabeledSystem, and MultiSystems. + + Notes + ----- + Ensure all plugins have been loaded before execuating this function! + """ + for method, formatcls in Format.get_from_methods().items(): + def get_func(ff): + # ff is not initized when defining from_format so cannot be polluted + def from_format(self, file_name, **kwargs): + return self.from_fmt_obj(ff(), file_name, **kwargs) + return from_format + + setattr(System, method, get_func(formatcls)) + setattr(LabeledSystem, method, get_func(formatcls)) + setattr(MultiSystems, method, get_func(formatcls)) + + for method, formatcls in Format.get_to_methods().items(): + def get_func(ff): + def to_format(self, *args, **kwargs): + return self.to_fmt_obj(ff(), *args, **kwargs) + return to_format + + setattr(System, method, get_func(formatcls)) + setattr(LabeledSystem, method, get_func(formatcls)) + setattr(MultiSystems, method, get_func(formatcls)) + +add_format_methods() + def check_System(data): keys={'atom_names','atom_numbs','cells','coords','orig','atom_types'} assert( isinstance(data,dict) ) diff --git a/plugin_example/README.md b/plugin_example/README.md new file mode 100644 index 000000000..0c5b8a326 --- /dev/null +++ b/plugin_example/README.md @@ -0,0 +1,22 @@ +# An example for plugin + +```sh +pip install . +``` + +```py +import dpdata +print(dpdata.System(12, fmt="random")) +``` + +The output is + +> Data Summary +> Unlabeled System +> ------------------- +> Frame Numbers : 12 +> Atom Numbers : 20 +> Element List : +> ------------------- +> X +> 20 \ No newline at end of file diff --git a/plugin_example/dpdata_random/__init__.py b/plugin_example/dpdata_random/__init__.py new file mode 100644 index 000000000..b9ce840d7 --- /dev/null +++ b/plugin_example/dpdata_random/__init__.py @@ -0,0 +1,24 @@ +from dpdata.format import Format +import numpy as np + +@Format.register("random") +class RandomFormat(Format): + def from_system(self, N, **kwargs): + return { + "atom_numbs": [20], + "atom_names": ['X'], + "atom_types": [0] * 20, + "cells": np.repeat(np.diag(np.diag(np.ones((3, 3))))[np.newaxis,...], N, axis=0) * 100., + "coords": np.random.rand(N, 20, 3) * 100., + } + + def from_labeled_system(self, N, **kwargs): + return { + "atom_numbs": [20], + "atom_names": ['X'], + "atom_types": [0] * 20, + "cells": np.repeat(np.diag(np.diag(np.ones((3, 3))))[np.newaxis,...], N, axis=0) * 100., + "coords": np.random.rand(N, 20, 3) * 100., + "energies": np.random.rand(N) * 100., + "forces": np.random.rand(N, 20, 3) * 100., + } \ No newline at end of file diff --git a/plugin_example/setup.py b/plugin_example/setup.py new file mode 100644 index 000000000..031afb60f --- /dev/null +++ b/plugin_example/setup.py @@ -0,0 +1,12 @@ +from setuptools import setup + +setup( + name="dpdata_random", + packages=['dpdata_random'], + entry_points={ + 'dpdata.plugins': [ + 'random=dpdata_random:RandomFormat' + ] + }, + install_requires=['dpdata', 'numpy'], +) diff --git a/tests/test_bond_order_system.py b/tests/test_bond_order_system.py index 086a90cad..c0c97a9a1 100644 --- a/tests/test_bond_order_system.py +++ b/tests/test_bond_order_system.py @@ -97,3 +97,6 @@ def test_sanitize_mol_origin(self): def tearDown(self): if os.path.exists("tests/.cache"): shutil.rmtree("tests/.cache") + +if __name__ == '__main__': + unittest.main() From b1025b479b718d3caff98d06792403b888d55478 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 3 Jul 2021 12:40:13 -0400 Subject: [PATCH 2/4] register functions automatically --- dpdata/plugins/abacus.py | 2 +- dpdata/plugins/amber.py | 2 -- dpdata/plugins/ase.py | 1 - dpdata/plugins/cp2k.py | 2 -- dpdata/plugins/deepmd.py | 7 +------ dpdata/plugins/fhi_aims.py | 3 +-- dpdata/plugins/gaussian.py | 2 -- dpdata/plugins/gromacs.py | 2 -- dpdata/plugins/lammps.py | 4 ---- dpdata/plugins/list.py | 1 - dpdata/plugins/pwmat.py | 4 +--- dpdata/plugins/pymatgen.py | 1 - dpdata/plugins/qe.py | 2 -- dpdata/plugins/rdkit.py | 6 ++---- dpdata/plugins/siesta.py | 1 - dpdata/plugins/vasp.py | 5 ----- dpdata/plugins/xyz.py | 2 +- dpdata/system.py | 9 +++++++++ 18 files changed, 16 insertions(+), 40 deletions(-) diff --git a/dpdata/plugins/abacus.py b/dpdata/plugins/abacus.py index e9998ce6a..ad37cdfa2 100644 --- a/dpdata/plugins/abacus.py +++ b/dpdata/plugins/abacus.py @@ -3,7 +3,7 @@ @Format.register("abacus/scf") -@Format.register_from("from_abacus_pw_scf") +@Format.register("abacus/pw/scf") class AbacusSCFFormat(Format): @Format.post("rot_lower_triangular") def from_labeled_system(self, file_name, **kwargs): diff --git a/dpdata/plugins/amber.py b/dpdata/plugins/amber.py index c7139a505..81e279138 100644 --- a/dpdata/plugins/amber.py +++ b/dpdata/plugins/amber.py @@ -4,7 +4,6 @@ @Format.register("amber/md") -@Format.register_from("from_amber_md") class AmberMDFormat(Format): def from_system(self, file_name=None, parm7_file=None, nc_file=None, use_element_symbols=None): # assume the prefix is the same if the spefic name is not given @@ -30,7 +29,6 @@ def from_labeled_system(self, file_name=None, parm7_file=None, nc_file=None, mdf @Format.register("sqm/out") -@Format.register_from("from_sqm_out") class SQMOutFormat(Format): def from_system(self, fname, **kwargs): ''' diff --git a/dpdata/plugins/ase.py b/dpdata/plugins/ase.py index 79ba9e0b6..4d0e0aedb 100644 --- a/dpdata/plugins/ase.py +++ b/dpdata/plugins/ase.py @@ -2,7 +2,6 @@ @Format.register("ase/structure") -@Format.register_to("to_ase_structure") class ASEStructureFormat(Format): def to_system(self, data, **kwargs): ''' diff --git a/dpdata/plugins/cp2k.py b/dpdata/plugins/cp2k.py index 4e2933f67..8787c7f09 100644 --- a/dpdata/plugins/cp2k.py +++ b/dpdata/plugins/cp2k.py @@ -5,7 +5,6 @@ @Format.register("cp2k/aimd_output") -@Format.register_from("from_cp2k_aimd_output") class CP2KAIMDOutputFormat(Format): def from_labeled_system(self, file_name, restart=False, **kwargs): xyz_file = sorted(glob.glob("{}/*pos*.xyz".format(file_name)))[0] @@ -14,7 +13,6 @@ def from_labeled_system(self, file_name, restart=False, **kwargs): @Format.register("cp2k/output") -@Format.register_from("from_cp2k_output") class CP2KOutputFormat(Format): def from_labeled_system(self, file_name, restart=False, **kwargs): data = {} diff --git a/dpdata/plugins/deepmd.py b/dpdata/plugins/deepmd.py index 0f7e8f30f..6910cc60a 100644 --- a/dpdata/plugins/deepmd.py +++ b/dpdata/plugins/deepmd.py @@ -6,8 +6,6 @@ @Format.register("deepmd") @Format.register("deepmd/raw") -@Format.register_from("from_deepmd_raw") -@Format.register_to("to_deepmd_raw") class DeePMDRawFormat(Format): def from_system(self, file_name, type_map=None, **kwargs): return dpdata.deepmd.raw.to_system_data(file_name, type_map=type_map, labels=False) @@ -24,10 +22,7 @@ def from_labeled_system(self, file_name, type_map, **kwargs): @Format.register("deepmd/npy") -@Format.register_from("from_deepmd_comp") -@Format.register_from("from_deepmd_npy") -@Format.register_to("to_deepmd_comp") -@Format.register_to("to_deepmd_npy") +@Format.register("deepmd/comp") class DeePMDCompFormat(Format): def from_system(self, file_name, type_map=None, **kwargs): return dpdata.deepmd.comp.to_system_data(file_name, type_map=type_map, labels=False) diff --git a/dpdata/plugins/fhi_aims.py b/dpdata/plugins/fhi_aims.py index e4b1b6239..96e000c6b 100644 --- a/dpdata/plugins/fhi_aims.py +++ b/dpdata/plugins/fhi_aims.py @@ -2,7 +2,7 @@ from dpdata.format import Format @Format.register("fhi_aims/md") -@Format.register_from("from_fhi_aims_output") +@Format.register("fhi_aims/output") class FhiMDFormat(Format): def from_labeled_system(self, file_name, md=True, begin = 0, step = 1, **kwargs): data = {} @@ -20,7 +20,6 @@ def from_labeled_system(self, file_name, md=True, begin = 0, step = 1, **kwargs) return data @Format.register("fhi_aims/scf") -@Format.register_from("from_fhi_aims_output") class FhiSCFFormat(Format): def from_labeled_system(self, file_name, **kwargs): data = {} diff --git a/dpdata/plugins/gaussian.py b/dpdata/plugins/gaussian.py index 95cd3d755..96e0993e3 100644 --- a/dpdata/plugins/gaussian.py +++ b/dpdata/plugins/gaussian.py @@ -3,7 +3,6 @@ @Format.register("gaussian/log") -@Format.register_from("from_gaussian_log") class GaussianLogFormat(Format): def from_labeled_system(self, file_name, md=False, **kwargs): try: @@ -17,7 +16,6 @@ def from_labeled_system(self, file_name, md=False, **kwargs): @Format.register("gaussian/md") -@Format.register_from("from_gaussian_md") class GaussianMDFormat(Format): def from_labeled_system(self, file_name, **kwargs): return GaussianLogFormat().from_labeled_system(file_name, md=True) diff --git a/dpdata/plugins/gromacs.py b/dpdata/plugins/gromacs.py index 6335c6d90..c45435836 100644 --- a/dpdata/plugins/gromacs.py +++ b/dpdata/plugins/gromacs.py @@ -4,8 +4,6 @@ @Format.register("gro") @Format.register("gromacs/gro") -@Format.register_from("from_gromacs_gro") -@Format.register_to("to_gromacs_gro") class PwmatOutputFormat(Format): def from_system(self, file_name, format_atom_name=True, **kwargs): """ diff --git a/dpdata/plugins/lammps.py b/dpdata/plugins/lammps.py index eb5dd3c8d..f1356bb1d 100644 --- a/dpdata/plugins/lammps.py +++ b/dpdata/plugins/lammps.py @@ -5,8 +5,6 @@ @Format.register("lmp") @Format.register("lammps/lmp") -@Format.register_from("from_lammps_lmp") -@Format.register_to("to_lammps_lmp") class LAMMPSLmpFormat(Format): @Format.post("shift_orig_zero") def from_system(self, file_name, type_map=None, **kwargs): @@ -35,8 +33,6 @@ def to_system(self, data, file_name, frame_idx=0, **kwargs): @Format.register("dump") @Format.register("lammps/dump") -@Format.register_from("from_lammps_dump") -@Format.register_to("to_lammps_dump") class LAMMPSDumpFormat(Format): @Format.post("shift_orig_zero") def from_system(self, diff --git a/dpdata/plugins/list.py b/dpdata/plugins/list.py index ffe0c5381..0eca2e13d 100644 --- a/dpdata/plugins/list.py +++ b/dpdata/plugins/list.py @@ -2,7 +2,6 @@ @Format.register("list") -@Format.register_to("to_list") class ListFormat(Format): def to_system(self, data, **kwargs): """ diff --git a/dpdata/plugins/pwmat.py b/dpdata/plugins/pwmat.py index 0d18e466e..7756e0c5c 100644 --- a/dpdata/plugins/pwmat.py +++ b/dpdata/plugins/pwmat.py @@ -8,7 +8,7 @@ @Format.register("mlmd") @Format.register("pwmat/movement") @Format.register("pwmat/mlmd") -@Format.register_from("from_pwmat_output") +@Format.register("pwmat/output") class PwmatOutputFormat(Format): @Format.post("rot_lower_triangular") def from_labeled_system(self, file_name, begin=0, step=1, **kwargs): @@ -37,8 +37,6 @@ def from_labeled_system(self, file_name, begin=0, step=1, **kwargs): @Format.register("final.config") @Format.register("pwmat/atom.config") @Format.register("pwmat/final.config") -@Format.register_from("from_pwmat_atomconfig") -@Format.register_to("to_pwmat_atomconfig") class PwmatAtomconfigFormat(Format): @Format.post("rot_lower_triangular") def from_system(self, file_name, **kwargs): diff --git a/dpdata/plugins/pymatgen.py b/dpdata/plugins/pymatgen.py index 8e8b63a9a..a047249e1 100644 --- a/dpdata/plugins/pymatgen.py +++ b/dpdata/plugins/pymatgen.py @@ -2,7 +2,6 @@ @Format.register("pymatgen/structure") -@Format.register_to("to_pymatgen_structure") class PyMatgenStructureFormat(Format): def to_system(self, data, **kwargs): """convert System to Pymatgen Structure obj diff --git a/dpdata/plugins/qe.py b/dpdata/plugins/qe.py index 17641ec8d..0e8149d00 100644 --- a/dpdata/plugins/qe.py +++ b/dpdata/plugins/qe.py @@ -4,7 +4,6 @@ from dpdata.format import Format @Format.register("qe/cp/traj") -@Format.register_from("from_qe_cp_traj") class QECPTrajFormat(Format): @Format.post("rot_lower_triangular") def from_system(self, file_name, begin = 0, step = 1, **kwargs): @@ -28,7 +27,6 @@ def from_labeled_system(self, file_name, begin = 0, step = 1, **kwargs): return data @Format.register("qe/pw/scf") -@Format.register_from("from_qe_pw_scf") class QECPTrajFormat(Format): @Format.post("rot_lower_triangular") def from_labeled_system(self, file_name, **kwargs): diff --git a/dpdata/plugins/rdkit.py b/dpdata/plugins/rdkit.py index bc19bf3a1..20e5696b2 100644 --- a/dpdata/plugins/rdkit.py +++ b/dpdata/plugins/rdkit.py @@ -7,8 +7,7 @@ @Format.register("mol") -@Format.register_from("from_mol_file") -@Format.register_to("to_mol_file") +@Format.register("mol_file") class MolFormat(Format): def from_bond_order_system(self, file_name, **kwargs): return rdkit.Chem.MolFromMolFile(file_name, sanitize=False, removeHs=False) @@ -20,8 +19,7 @@ def to_bond_order_system(self, mol, data, file_name, frame_idx=0, **kwargs): @Format.register("sdf") -@Format.register_from("from_sdf_file") -@Format.register_to("to_sdf_file") +@Format.register("sdf_file") class SdfFormat(Format): def from_bond_order_system(self, file_name, **kwargs): ''' diff --git a/dpdata/plugins/siesta.py b/dpdata/plugins/siesta.py index c6bb55399..a601e1324 100644 --- a/dpdata/plugins/siesta.py +++ b/dpdata/plugins/siesta.py @@ -4,7 +4,6 @@ @Format.register("siesta/output") -@Format.register_from("from_siesta_output") class SiestaOutputFormat(Format): def from_system(self, file_name, **kwargs): data = {} diff --git a/dpdata/plugins/vasp.py b/dpdata/plugins/vasp.py index fca08f15d..01aad2026 100644 --- a/dpdata/plugins/vasp.py +++ b/dpdata/plugins/vasp.py @@ -9,8 +9,6 @@ @Format.register("contcar") @Format.register("vasp/poscar") @Format.register("vasp/contcar") -@Format.register_from("from_vasp_poscar") -@Format.register_to("to_vasp_poscar") class VASPPoscarFormat(Format): @Format.post("rot_lower_triangular") def from_system(self, file_name, **kwargs): @@ -35,7 +33,6 @@ def to_system(self, data, file_name, frame_idx=0, **kwargs): @Format.register("vasp/string") -@Format.register_to("to_vasp_string") class VASPStringFormat(Format): def to_system(self, data, frame_idx=0, **kwargs): """ @@ -53,7 +50,6 @@ def to_system(self, data, frame_idx=0, **kwargs): # rotate the system to lammps convention @Format.register("outcar") @Format.register("vasp/outcar") -@Format.register_from("from_vasp_outcar") class VASPOutcarFormat(Format): @Format.post("rot_lower_triangular") def from_labeled_system(self, file_name, begin=0, step=1, **kwargs): @@ -81,7 +77,6 @@ def from_labeled_system(self, file_name, begin=0, step=1, **kwargs): # rotate the system to lammps convention @Format.register("xml") @Format.register("vasp/xml") -@Format.register_from("from_vasp_xml") class VASPXMLFormat(Format): @Format.post("rot_lower_triangular") def from_labeled_system(self, file_name, begin=0, step=1, **kwargs): diff --git a/dpdata/plugins/xyz.py b/dpdata/plugins/xyz.py index a3fa7b27c..b4e77c9dd 100644 --- a/dpdata/plugins/xyz.py +++ b/dpdata/plugins/xyz.py @@ -3,7 +3,7 @@ @Format.register("quip/gap/xyz") -@Format.register_from("from_quip_gap_xyz_file") +@Format.register("quip/gap/xyz_file") class DeePMDRawFormat(Format): def from_labeled_system(self, data, **kwargs): return data diff --git a/dpdata/system.py b/dpdata/system.py index 5ca0f24bd..f2222e4eb 100644 --- a/dpdata/system.py +++ b/dpdata/system.py @@ -1221,6 +1221,15 @@ def add_format_methods(): ----- Ensure all plugins have been loaded before execuating this function! """ + # automatically register from/to functions for formats + # for example, deepmd/npy will be registered as from_deepmd_npy and to_deepmd_npy + for key, formatcls in Format.get_formats().items(): + formattedkey = key.replace("/", "_").replace(".", "") + from_func_name = "from_" + formattedkey + to_func_name = "to_" + formattedkey + Format.register_from(from_func_name)(formatcls) + Format.register_to(to_func_name)(formatcls) + for method, formatcls in Format.get_from_methods().items(): def get_func(ff): # ff is not initized when defining from_format so cannot be polluted From 3967210edc04dbe04419969800748400ba8d9db7 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 3 Jul 2021 16:13:26 -0400 Subject: [PATCH 3/4] add dpdata/plugins to packages --- setup.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 6de24481d..8fb19676c 100644 --- a/setup.py +++ b/setup.py @@ -39,7 +39,8 @@ 'dpdata/fhi_aims', 'dpdata/gromacs', 'dpdata/abacus', - 'dpdata/rdkit' + 'dpdata/rdkit', + 'dpdata/plugins', ], package_data={'dpdata':['*.json']}, classifiers=[ From 4ca370691ed7b7403c1a3dce83c3616b5a90a300 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 4 Jul 2021 05:40:54 -0400 Subject: [PATCH 4/4] fix typo, add docs --- README.md | 11 +++++++++++ dpdata/plugins/gromacs.py | 2 +- dpdata/plugins/xyz.py | 2 +- plugin_example/README.md | 20 +++++++++++--------- 4 files changed, 24 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 66333f824..e62802992 100644 --- a/README.md +++ b/README.md @@ -252,3 +252,14 @@ print(syst.get_charge()) # return the total charge of the system ``` If a valence of 3 is detected on carbon, the formal charge will be assigned to -1. Because for most cases (in alkynyl anion, isonitrile, cyclopentadienyl anion), the formal charge on 3-valence carbon is -1, and this is also consisent with the 8-electron rule. + +# Plugins + +One can follow [a simple example](plugin_example/) to add their own format by creating and installing plugins. It's crirical to add the [Format](dpdata/format.py) class to `entry_points['dpdata.plugins']` in `setup.py`: +```py + entry_points={ + 'dpdata.plugins': [ + 'random=dpdata_random:RandomFormat' + ] + }, +``` diff --git a/dpdata/plugins/gromacs.py b/dpdata/plugins/gromacs.py index c45435836..a061918ba 100644 --- a/dpdata/plugins/gromacs.py +++ b/dpdata/plugins/gromacs.py @@ -4,7 +4,7 @@ @Format.register("gro") @Format.register("gromacs/gro") -class PwmatOutputFormat(Format): +class GromacsGroFormat(Format): def from_system(self, file_name, format_atom_name=True, **kwargs): """ Load gromacs .gro file diff --git a/dpdata/plugins/xyz.py b/dpdata/plugins/xyz.py index b4e77c9dd..71bf92683 100644 --- a/dpdata/plugins/xyz.py +++ b/dpdata/plugins/xyz.py @@ -4,7 +4,7 @@ @Format.register("quip/gap/xyz") @Format.register("quip/gap/xyz_file") -class DeePMDRawFormat(Format): +class QuipGapXYZFormat(Format): def from_labeled_system(self, data, **kwargs): return data diff --git a/plugin_example/README.md b/plugin_example/README.md index 0c5b8a326..10aadf040 100644 --- a/plugin_example/README.md +++ b/plugin_example/README.md @@ -11,12 +11,14 @@ print(dpdata.System(12, fmt="random")) The output is -> Data Summary -> Unlabeled System -> ------------------- -> Frame Numbers : 12 -> Atom Numbers : 20 -> Element List : -> ------------------- -> X -> 20 \ No newline at end of file +``` +Data Summary +Unlabeled System +------------------- +Frame Numbers : 12 +Atom Numbers : 20 +Element List : +------------------- +X +20 +``` \ No newline at end of file