diff --git a/dpdata/system.py b/dpdata/system.py index bba7446f0..f5e1ffb3e 100644 --- a/dpdata/system.py +++ b/dpdata/system.py @@ -40,6 +40,10 @@ def decorator(func): return func return decorator + def __add__(self, other): + self.funcs.update(other.funcs) + return self + class System (MSONable) : ''' @@ -309,9 +313,11 @@ def sub_system(self, f_idx) : tmp = System() for ii in ['atom_numbs', 'atom_names', 'atom_types', 'orig'] : tmp.data[ii] = self.data[ii] + tmp.data['cells'] = self.data['cells'][f_idx].reshape(-1, 3, 3) tmp.data['coords'] = self.data['coords'][f_idx].reshape(-1, self.data['coords'].shape[1], 3) tmp.data['nopbc'] = self.nopbc + return tmp @@ -516,16 +522,15 @@ def to_ase_structure(self): convert System to ASE Atom obj ''' + from ase import Atoms + structures=[] - try: - from ase import Atoms - except: - raise ImportError('No module ase.Atoms') 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") @@ -1045,6 +1050,7 @@ def __init__ (self, self.apply_type_map(type_map) register_from_funcs = Register() + register_to_funcs = System.register_to_funcs + Register() def __repr__(self): return self.__str__() @@ -1338,6 +1344,37 @@ 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): """