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)