def compute_distances(traj, atom_pairs, periodic=True, opt=True):
"""Compute the distances between pairs of atoms in each frame.
Parameters
----------
traj : Trajectory
An mtraj trajectory.
atom_pairs : np.ndarray, shape=(num_pairs, 2), dtype=int
Each row gives the indices of two atoms involved in the interaction.
periodic : bool, default=True
If `periodic` is True and the trajectory contains unitcell
information, we will compute distances under the minimum image
convention.
opt : bool, default=True
Use an optimized native library to calculate distances. Our optimized
SSE minimum image convention calculation implementation is over 1000x
faster than the naive numpy implementation.
Returns
-------
distances : np.ndarray, shape=(n_frames, num_pairs), dtype=float
The distance, in each frame, between each pair of atoms.
"""
xyz = ensure_type(traj.xyz,
dtype=np.float32,
ndim=3,
name='traj.xyz',
shape=(None, None, 3),
warn_on_cast=False)
pairs = ensure_type(np.asarray(atom_pairs),
dtype=np.int32,
ndim=2,
name='atom_pairs',
shape=(None, 2),
warn_on_cast=False)
if not np.all(np.logical_and(pairs < traj.n_atoms, pairs >= 0)):
raise ValueError('atom_pairs must be between 0 and %d' % traj.n_atoms)
if periodic is True and traj._have_unitcell:
box = ensure_type(traj.unitcell_vectors,
dtype=np.float32,
ndim=3,
name='unitcell_vectors',
shape=(len(xyz), 3, 3))
if opt and _geometry._processor_supports_sse41():
out = np.empty((xyz.shape[0], pairs.shape[0]), dtype=np.float32)
_geometry._dist_mic(xyz, pairs, box, out)
return out
else:
return _distance_mic(xyz, pairs, box)
# either there are no unitcell vectors or they dont want to use them
if opt and _geometry._processor_supports_sse41():
out = np.empty((xyz.shape[0], pairs.shape[0]), dtype=np.float32)
_geometry._dist(xyz, pairs, out)
return out
else:
return _distance(xyz, pairs)
def _compute_distances_chunk(xyz,
pairs,
box=None,
periodic=True,
opt=True,
orthogonal=False):
"""Compute distances for a single chunk
Parameters
----------
xyz : ndarray of shape (any, any, 3)
The xyz coordinates of the chunk
pairs : array of shape (any, 2)
The indices for which to compute an angle
box : ndarray of shape (any, 3, 3)
The box vectors of the chunk
periodic : bool
Wether to use the periodc boundary during the calculation.
opt : bool, default=True
Use an optimized native library to calculate distances. MDTraj's
optimized SSE angle calculation implementation is 10-20x faster than
the (itself optimized) numpy implementation.
orthogonal : bool or da.bool
Wether all angles are close to 90 degrees
"""
xyz = ensure_type(
xyz,
dtype=np.float32,
ndim=3,
name="xyz",
shape=(None, None, 3),
warn_on_cast=False,
cast_da_to_np=True,
)
# Cast orthogonal to bool (incase we get a delayed bool)
orthogonal = bool(orthogonal)
if periodic and box is not None:
if opt:
out = np.empty((xyz.shape[0], pairs.shape[0]), dtype=np.float32)
_geometry._dist_mic(xyz, pairs,
box.transpose(0, 2, 1).copy(), out, orthogonal)
return out
else:
return _distance_mic(xyz, pairs, box.transpose(0, 2, 1),
orthogonal)
# Either there are no unitcell vectors or they dont want to use them
if opt:
out = np.empty((xyz.shape[0], pairs.shape[0]), dtype=np.float32)
_geometry._dist(xyz, pairs, out)
return out
else:
return _distance(xyz, pairs)
def compute_distances(traj, atom_pairs, periodic=True, opt=True):
"""Compute the distances between pairs of atoms in each frame.
Parameters
----------
traj : Trajectory
An mtraj trajectory.
atom_pairs : np.ndarray, shape=(num_pairs, 2), dtype=int
Each row gives the indices of two atoms involved in the interaction.
periodic : bool, default=True
If `periodic` is True and the trajectory contains unitcell
information, we will compute distances under the minimum image
convention.
opt : bool, default=True
Use an optimized native library to calculate distances. Our optimized
SSE minimum image convention calculation implementation is over 1000x
faster than the naive numpy implementation.
Returns
-------
distances : np.ndarray, shape=(n_frames, num_pairs), dtype=float
The distance, in each frame, between each pair of atoms.
"""
xyz = ensure_type(traj.xyz, dtype=np.float32, ndim=3, name='traj.xyz', shape=(None, None, 3), warn_on_cast=False)
pairs = ensure_type(atom_pairs, dtype=np.int32, ndim=2, name='atom_pairs', shape=(None, 2), warn_on_cast=False)
if not np.all(np.logical_and(pairs < traj.n_atoms, pairs >= 0)):
raise ValueError('atom_pairs must be between 0 and %d' % traj.n_atoms)
if len(pairs) == 0:
return np.zeros((len(xyz), 0), dtype=np.float32)
if periodic and traj._have_unitcell:
box = ensure_type(traj.unitcell_vectors, dtype=np.float32, ndim=3, name='unitcell_vectors', shape=(len(xyz), 3, 3),
warn_on_cast=False)
orthogonal = np.allclose(traj.unitcell_angles, 90)
if opt:
out = np.empty((xyz.shape[0], pairs.shape[0]), dtype=np.float32)
_geometry._dist_mic(xyz, pairs, box.transpose(0, 2, 1).copy(), out, orthogonal)
return out
else:
return _distance_mic(xyz, pairs, box.transpose(0, 2, 1), orthogonal)
# either there are no unitcell vectors or they dont want to use them
if opt:
out = np.empty((xyz.shape[0], pairs.shape[0]), dtype=np.float32)
_geometry._dist(xyz, pairs, out)
return out
else:
return _distance(xyz, pairs)
def compute_distances(traj, atom_pairs, periodic=True, opt=True):
"""Compute the distances between pairs of atoms in each frame.
Parameters
----------
traj : Trajectory
An mtraj trajectory.
atom_pairs : np.ndarray, shape=(num_pairs, 2), dtype=int
Each row gives the indices of two atoms involved in the interaction.
periodic : bool, default=True
If `periodic` is True and the trajectory contains unitcell
information, we will compute distances under the minimum image
convention.
opt : bool, default=True
Use an optimized native library to calculate distances. Our optimized
SSE minimum image convention calculation implementation is over 1000x
faster than the naive numpy implementation.
Returns
-------
distances : np.ndarray, shape=(n_frames, num_pairs), dtype=float
The distance, in each frame, between each pair of atoms.
"""
xyz = ensure_type(traj.xyz, dtype=np.float32, ndim=3, name='taj.xyz', shape=(None, None, 3), warn_on_cast=False)
pairs = ensure_type(np.asarray(atom_pairs), dtype=np.int32, ndim=2, name='atom_pairs', shape=(None, 2), warn_on_cast=False)
if periodic is True and traj._have_unitcell:
box = ensure_type(traj.unitcell_vectors, dtype=np.float32, ndim=3, name='unitcell_vectors', shape=(len(xyz), 3, 3))
if opt:
out = np.empty((xyz.shape[0], pairs.shape[0]), dtype=np.float32)
_geometry._dist_mic(xyz, pairs, box, out)
return out
else:
return _distance_mic(xyz, pairs, box)
# either there are no unitcell vectors or they dont want to use them
if opt:
out = np.empty((xyz.shape[0], pairs.shape[0]), dtype=np.float32)
_geometry._dist(xyz, pairs, out)
return out
else:
return _distance(xyz, pairs)
