From aad2b4a2108c637168308d57b7b55d258eee1e9a Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 22 May 2022 23:35:16 -0400 Subject: [PATCH] support converting ase.Atoms to System and LabeledSystem #190 added `from_labeled_system`, but it just passes a dict and does nothing else. This commit rewrites it. --- docs/conf.py | 3 + dpdata/plugins/ase.py | 126 +++++++++++++++++++++++++++++++----------- tests/test_to_ase.py | 30 ++++++++-- 3 files changed, 121 insertions(+), 38 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index dbe2c0bf8..d82fa6a48 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -183,4 +183,7 @@ def setup(app): intersphinx_mapping = { "numpy": ("https://docs.scipy.org/doc/numpy/", None), "python": ("https://docs.python.org/", None), + "ase": ("https://wiki.fysik.dtu.dk/ase/", None), + "monty": ("https://guide.materialsvirtuallab.org/monty/", None), + "h5py": ("https://docs.h5py.org/en/stable/", None), } diff --git a/dpdata/plugins/ase.py b/dpdata/plugins/ase.py index 891b2d90a..957aac568 100644 --- a/dpdata/plugins/ase.py +++ b/dpdata/plugins/ase.py @@ -18,41 +18,101 @@ class ASEStructureFormat(Format): automatic detection fails. """ - def from_labeled_system(self, data, **kwargs): - return data - - def from_multi_systems(self, file_name, begin=None, end=None, step=None, ase_fmt=None, **kwargs): + def from_system(self, atoms: "ase.Atoms", **kwargs) -> dict: + """Convert ase.Atoms to a System. + + Parameters + ---------- + atoms : ase.Atoms + an ASE Atoms, containing a structure + + Returns + ------- + dict + data dict + """ + symbols = atoms.get_chemical_symbols() + atom_names = list(set(symbols)) + atom_numbs = [symbols.count(symbol) for symbol in atom_names] + atom_types = np.array([atom_names.index(symbol) for symbol in symbols]).astype(int) + cells = atoms.cell[:] + coords = atoms.get_positions() + info_dict = { + 'atom_names': atom_names, + 'atom_numbs': atom_numbs, + 'atom_types': atom_types, + 'cells': np.array([cells]).astype('float32'), + 'coords': np.array([coords]).astype('float32'), + 'orig': [0,0,0], + } + return info_dict + + def from_labeled_system(self, atoms: "ase.Atoms", **kwargs) -> dict: + """Convert ase.Atoms to a LabeledSystem. Energies and forces + are calculated by the calculator. + + Parameters + ---------- + atoms : ase.Atoms + an ASE Atoms, containing a structure + + Returns + ------- + dict + data dict + + Raises + ------ + RuntimeError + ASE will raise RuntimeError if the atoms does not + have a calculator + """ + info_dict = self.from_system(atoms) + try: + energies = atoms.get_potential_energy(force_consistent=True) + except PropertyNotImplementedError: + energies = atoms.get_potential_energy() + forces = atoms.get_forces() + info_dict = { + ** info_dict, + 'energies': np.array([energies]).astype('float32'), + 'forces': np.array([forces]).astype('float32'), + } + try: + stress = atoms.get_stress(False) + except PropertyNotImplementedError: + pass + else: + virials = np.array([-atoms.get_volume() * stress]).astype('float32') + info_dict['virials'] = virials + return info_dict + + def from_multi_systems(self, file_name: str, begin: int=None, end: int=None, step: int=None, ase_fmt: str=None, **kwargs) -> "ase.Atoms": + """Convert a ASE supported file to ASE Atoms. + + It will finally be converted to MultiSystems. + + Parameters + ---------- + file_name : str + path to file + begin : int, optional + begin frame index + end : int, optional + end frame index + step : int, optional + frame index step + ase_fmt : str, optional + ASE format. See the ASE documentation about supported formats + + Yields + ------ + ase.Atoms + ASE atoms in the file + """ frames = ase.io.read(file_name, format=ase_fmt, index=slice(begin, end, step)) for atoms in frames: - symbols = atoms.get_chemical_symbols() - atom_names = list(set(symbols)) - atom_numbs = [symbols.count(symbol) for symbol in atom_names] - atom_types = np.array([atom_names.index(symbol) for symbol in symbols]).astype(int) - - cells = atoms.cell[:] - coords = atoms.get_positions() - try: - energies = atoms.get_potential_energy(force_consistent=True) - except PropertyNotImplementedError: - energies = atoms.get_potential_energy() - forces = atoms.get_forces() - info_dict = { - 'atom_names': atom_names, - 'atom_numbs': atom_numbs, - 'atom_types': atom_types, - 'cells': np.array([cells]).astype('float32'), - 'coords': np.array([coords]).astype('float32'), - 'energies': np.array([energies]).astype('float32'), - 'forces': np.array([forces]).astype('float32'), - 'orig': [0,0,0], - } - try: - stress = atoms.get_stress(False) - virials = np.array([-atoms.get_volume() * stress]).astype('float32') - info_dict['virials'] = virials - except PropertyNotImplementedError: - pass - yield info_dict + yield atoms def to_system(self, data, **kwargs): ''' diff --git a/tests/test_to_ase.py b/tests/test_to_ase.py index 4ea870c79..d1c42b8b4 100644 --- a/tests/test_to_ase.py +++ b/tests/test_to_ase.py @@ -4,11 +4,13 @@ from context import dpdata from comp_sys import CompSys, IsPBC try: - from ase import Atoms - from ase.io import write - exist_module=True -except Exception: - exist_module=False + from ase import Atoms + from ase.io import write +except ModuleNotFoundError: + exist_module=False +else: + exist_module=True + @unittest.skipIf(not exist_module,"skip test_ase") class TestASE(unittest.TestCase, CompSys, IsPBC): @@ -24,6 +26,24 @@ def setUp(self): self.f_places = 6 self.v_places = 6 + +@unittest.skipIf(not exist_module, "skip test_ase") +class TestFromASE(unittest.TestCase, CompSys, IsPBC): + """Test ASEStructureFormat.from_system""" + def setUp(self): + system_1 = dpdata.System() + system_1.from_lammps_lmp(os.path.join('poscars', 'conf.lmp'), type_map = ['O', 'H']) + atoms = system_1.to_ase_structure()[0] + self.system_1 = system_1 + self.system_2 = dpdata.System(atoms, fmt="ase/structure") + # assign the same type_map + self.system_2.sort_atom_names(type_map=self.system_1.get_atom_names()) + self.places = 6 + self.e_places = 6 + self.f_places = 6 + self.v_places = 6 + + if __name__ == '__main__': unittest.main()