def atom_distances(rst, atom_pairs):
"""Finds the distance between the each of the atom pairs in the
given LAMMPS dump file.
:param rst: A file in the LAMMPS dump format.
:param atom_pairs: Zero or more pairs of atom IDs to compare.
:returns: Nested dicts keyed by time step, then pair, with the distance as the value.
"""
results = OrderedDict()
flat_ids = set(itertools.chain.from_iterable(atom_pairs))
tstep_atoms, tstep_box = find_atom_data(rst, flat_ids)
for tstep, atoms in tstep_atoms.items():
pair_dist = OrderedDict({FILENAME: os.path.basename(rst)})
for pair in atom_pairs:
try:
row1 = atoms[pair[0]]
row2 = atoms[pair[1]]
pair_dist[pair] = pbc_dist(row1[-3:], row2[-3:],
tstep_box[tstep])
except KeyError as e:
warning(MISSING_TSTEP_ATOM_MSG.format(rst, tstep, e))
return
results[tstep] = pair_dist
return results
def atom_distances(rst, atom_pairs):
"""Finds the distance between the each of the atom pairs in the
given LAMMPS dump file.
:param rst: A file in the LAMMPS dump format.
:param atom_pairs: Zero or more pairs of atom IDs to compare.
:returns: Nested dicts keyed by time step, then pair, with the distance as the value.
"""
results = OrderedDict()
flat_ids = set(itertools.chain.from_iterable(atom_pairs))
tstep_atoms = find_atom_data(rst, flat_ids)
for tstep, atoms in tstep_atoms.items():
pair_dist = OrderedDict()
for pair in atom_pairs:
try:
row1 = atoms[pair[0]]
row2 = atoms[pair[1]]
pair_dist[pair] = xyz_distance(row1[-3:], row2[-3:])
except KeyError as e:
raise InvalidDataError(MISSING_TSTEP_ATOM_MSG.format(rst, tstep, e))
results[tstep] = pair_dist
return results
def testMissingAtoms(self):
with self.assertRaises(InvalidDataError):
find_atom_data(DUMP_PATH, {-97})
def testGood(self):
lam_atoms, lam_box = find_atom_data(DUMP_PATH, {4167, 17467})
for box_dim in lam_box.values():
self.assertTrue(np.allclose(box_dim, LAM_BOX))
self.assertEqual(LAM_ATOMS, lam_atoms)
def testMissingAtoms(self):
with self.assertRaises(InvalidDataError):
find_atom_data(DUMP_PATH, {-97})
def testGood(self):
self.assertEqual(LAM_ATOMS, find_atom_data(DUMP_PATH, {4167, 17467}))
