def compute_displacements(traj, atom_pairs, periodic=True, opt=True): """Compute the displacement vector between pairs of atoms in each frame of a trajectory. Parameters ---------- traj : Trajectory Trajectory to compute distances in atom_pairs : np.ndarray, shape[num_pairs, 2], dtype=int Each row gives the indices of two atoms. 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 minimum image convention calculation implementation is over 1000x faster than the naive numpy implementation. Returns ------- displacements : np.ndarray, shape=[n_frames, n_pairs, 3], dtype=float32 The displacememt vector, 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)) pairs = ensure_type(np.asarray(atom_pairs), dtype=np.int32, ndim=2, name='atom_pairs', shape=(None, 2)) 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], 3), dtype=np.float32) _geometry._dist_mic_displacement(xyz, pairs, box, out) return out else: return _displacement_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], 3), dtype=np.float32) _geometry._dist_displacement(xyz, pairs, out) return out return _displacement(xyz, pairs)
def _compute_displacements_chunk(xyz, pairs, box=None, periodic=True, opt=True, orthogonal=False): """Compute displacements 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 """ # Cast orthogonal to a bool, just incase we got a delayed object orthogonal = bool(orthogonal) xyz = ensure_type( xyz, dtype=np.float32, ndim=3, name="xyz", shape=(None, None, 3), warn_on_cast=False, cast_da_to_np=True, ) if periodic and box is not None: if opt: out = np.empty((xyz.shape[0], pairs.shape[0], 3), dtype=np.float32) _geometry._dist_mic_displacement(xyz, pairs, box.transpose(0, 2, 1).copy(), out, orthogonal) return out else: return _displacement_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], 3), dtype=np.float32) _geometry._dist_displacement(xyz, pairs, out) return out else: return _displacement(xyz, pairs)
def compute_displacements(traj, atom_pairs, periodic=True, opt=True): """Compute the displacement vector between pairs of atoms in each frame of a trajectory. Parameters ---------- traj : Trajectory Trajectory to compute distances in atom_pairs : np.ndarray, shape[num_pairs, 2], dtype=int Each row gives the indices of two atoms. 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 minimum image convention calculation implementation is over 1000x faster than the naive numpy implementation. Returns ------- displacements : np.ndarray, shape=[n_frames, n_pairs, 3], dtype=float32 The displacememt vector, 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 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], 3), dtype=np.float32) _geometry._dist_mic_displacement(xyz, pairs, box.transpose(0, 2, 1).copy(), out, orthogonal) return out else: return _displacement_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], 3), dtype=np.float32) _geometry._dist_displacement(xyz, pairs, out) return out return _displacement(xyz, pairs)