diff --git a/dpdata/lammps/dump.py b/dpdata/lammps/dump.py index 7136d3e63..e40dc400b 100644 --- a/dpdata/lammps/dump.py +++ b/dpdata/lammps/dump.py @@ -5,6 +5,11 @@ lib_path = os.path.dirname(os.path.realpath(__file__)) sys.path.append(lib_path) import lmp +import warnings +class UnwrapWarning(UserWarning): + pass +warnings.simplefilter('once', UnwrapWarning) + def _get_block (lines, key) : for idx in range(len(lines)) : @@ -59,16 +64,17 @@ def get_coordtype_and_scalefactor(keys): key_su = ['xsu','ysu','zsu'] #scaled and unfolded,sf = lattice parameter lmp_coor_type = [key_pc,key_uc,key_s,key_su] sf = [0,0,1,1] + uw = [0,1,0,1] # unwraped or not for k in range(4): if all(i in keys for i in lmp_coor_type[k]): - return lmp_coor_type[k],sf[k] + return lmp_coor_type[k], sf[k], uw[k] -def safe_get_posi(lines,cell,orig=np.zeros(3)) : +def safe_get_posi(lines,cell,orig=np.zeros(3), unwrap=False) : blk, head = _get_block(lines, 'ATOMS') keys = head.split() coord_tp_and_sf = get_coordtype_and_scalefactor(keys) assert coord_tp_and_sf is not None, 'Dump file does not contain atomic coordinates!' - coordtype, sf = coord_tp_and_sf + coordtype, sf, uw = coord_tp_and_sf id_idx = keys.index('id') - 2 xidx = keys.index(coordtype[0])-2 yidx = keys.index(coordtype[1])-2 @@ -81,8 +87,13 @@ def safe_get_posi(lines,cell,orig=np.zeros(3)) : posis.sort() posis = np.array(posis)[:,1:4] if not sf: - posis = (posis-orig)@np.linalg.inv(cell)# Convert to scaled coordinates for unscaled coordinates - return (posis%1)@cell # Convert scaled coordinates back to Cartesien coordinates + posis = (posis - orig) @ np.linalg.inv(cell) # Convert to scaled coordinates for unscaled coordinates + if uw and unwrap: + return posis @ cell # convert scaled coordinates back to Cartesien coordinates unwrap at the periodic boundaries + else: + if uw and not unwrap: + warnings.warn(message='Your dump file contains unwrapped coordinates, but you did not specify unwrapping (unwrap = True). The default is wrapping at periodic boundaries (unwrap = False).\n',category=UnwrapWarning) + return (posis % 1) @ cell # Convert scaled coordinates back to Cartesien coordinates with wraping at periodic boundary conditions def get_dumpbox(lines) : blk, h = _get_block(lines, 'BOX BOUNDS') @@ -145,9 +156,9 @@ def load_file(fname, begin = 0, step = 1) : cc += 1 if cc >= begin and (cc - begin) % step == 0 : buff.append(line) - -def system_data(lines, type_map = None, type_idx_zero = True) : + +def system_data(lines, type_map = None, type_idx_zero = True, unwrap=False) : array_lines = split_traj(lines) lines = array_lines[0] system = {} @@ -166,12 +177,12 @@ def system_data(lines, type_map = None, type_idx_zero = True) : system['cells'] = [np.array(cell)] natoms = sum(system['atom_numbs']) system['atom_types'] = get_atype(lines, type_idx_zero = type_idx_zero) - system['coords'] = [safe_get_posi(lines, cell, np.array(orig))] + system['coords'] = [safe_get_posi(lines, cell, np.array(orig), unwrap)] for ii in range(1, len(array_lines)) : bounds, tilt = get_dumpbox(array_lines[ii]) orig, cell = dumpbox2box(bounds, tilt) system['cells'].append(cell) - system['coords'].append(safe_get_posi(array_lines[ii], cell, np.array(orig))) + system['coords'].append(safe_get_posi(array_lines[ii], cell, np.array(orig), unwrap)) system['cells'] = np.array(system['cells']) system['coords'] = np.array(system['coords']) return system @@ -181,7 +192,7 @@ def split_traj(dump_lines) : marks = [] for idx,ii in enumerate(dump_lines) : if 'ITEM: TIMESTEP' in ii : - marks.append(idx) + marks.append(idx) if len(marks) == 0 : return None elif len(marks) == 1 : @@ -191,9 +202,9 @@ def split_traj(dump_lines) : ret = [] for ii in marks : ret.append(dump_lines[ii:ii+block_size]) - # for ii in range(len(marks)-1): + # for ii in range(len(marks)-1): # assert(marks[ii+1] - marks[ii] == block_size) - return ret + return ret return None diff --git a/dpdata/plugins/lammps.py b/dpdata/plugins/lammps.py index f1356bb1d..8f9962967 100644 --- a/dpdata/plugins/lammps.py +++ b/dpdata/plugins/lammps.py @@ -40,6 +40,7 @@ def from_system(self, type_map=None, begin=0, step=1, + unwrap=False, **kwargs): lines = dpdata.lammps.dump.load_file(file_name, begin=begin, step=step) - return dpdata.lammps.dump.system_data(lines, type_map) + return dpdata.lammps.dump.system_data(lines, type_map, unwrap=unwrap) diff --git a/tests/poscars/poscar_ref_oh.py b/tests/poscars/poscar_ref_oh.py index 23bab644f..9b12c1511 100644 --- a/tests/poscars/poscar_ref_oh.py +++ b/tests/poscars/poscar_ref_oh.py @@ -27,8 +27,12 @@ def test_cell(self): places = 6, msg = 'cell[%d][%d] failed' % (ii,jj)) - def test_frame(self): - ovito_posis = np.array([[0, 0, 0], + def test_frame(self): + if hasattr(self, "unwrap") and self.unwrap is True: + ovito_posis = np.array([[5.0739861, 2.7916155, 2.2254033], + [6.3361717, 3.4934183, 2.7767918]]) + else: + ovito_posis = np.array([[0, 0, 0], [1.2621856, 0.7018028, 0.5513885]]) for ii in range(2) : for jj in range(3) : diff --git a/tests/test_lammps_dump_unfold.py b/tests/test_lammps_dump_unfold.py index 4ffd5f730..68aa2c8b7 100644 --- a/tests/test_lammps_dump_unfold.py +++ b/tests/test_lammps_dump_unfold.py @@ -21,6 +21,16 @@ def test_nframes (self) : self.assertEqual(self.tmp_system.get_nframes(), 2) +class TestDumpUnwrap(unittest.TestCase, TestPOSCARoh): + def setUp(self): + self.unwrap = True + self.system = dpdata.System( + os.path.join('poscars', 'conf_unfold.dump'), + type_map=['O', 'H'], + unwrap=self.unwrap, + ) + + if __name__ == '__main__': unittest.main()