def compute_angles(traj, angle_indices, periodic=True, opt=True): """Compute the bond angles between the supplied triplets of indices in each frame of a trajectory. Parameters ---------- traj : Trajectory An mdtraj trajectory. angle_indices : np.ndarray, shape=(num_angles, 3), dtype=int Each row gives the indices of three atoms which together make an angle. periodic : bool, default=True If `periodic` is True and the trajectory contains unitcell information, we will treat angles that cross periodic images using the minimum image convention. opt : bool, default=True Use an optimized native library to calculate distances. Our optimized SSE angle calculation implementation is 10-20x faster than the (itself optimized) numpy implementation. Returns ------- angles : np.ndarray, shape=[n_frames, n_angles], dtype=float The angles are in radians """ xyz = ensure_type(traj.xyz, dtype=np.float32, ndim=3, name='traj.xyz', shape=(None, None, 3), warn_on_cast=False) triplets = ensure_type(angle_indices, dtype=np.int32, ndim=2, name='angle_indices', shape=(None, 3), warn_on_cast=False) if not np.all(np.logical_and(triplets < traj.n_atoms, triplets >= 0)): raise ValueError('angle_indices must be between 0 and %d' % traj.n_atoms) if len(triplets) == 0: return np.zeros((len(xyz), 0), dtype=np.float32) if periodic and traj._have_unitcell: if opt and not np.allclose(traj.unitcell_angles, 90): warnings.warn('Optimized angle calculation does not work for non-orthorhombic ' 'unit cells and periodic boundary conditions. Falling back to much ' 'slower pure-Python implementation. Set periodic=False or opt=False ' 'to disable this message.') opt = False out = np.zeros((xyz.shape[0], triplets.shape[0]), dtype=np.float32) 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: _geometry._angle_mic(xyz, triplets, box, out) return out else: _angle(traj, triplets, periodic, out) return out if opt: _geometry._angle(xyz, triplets, out) else: _angle(traj, triplets, periodic, out) return out
def _compute_angles_chunk(xyz, triplets, box=None, periodic=True, opt=True, orthogonal=False): """Compute the angles for a single chunk Parameters ---------- xyz : ndarray of shape (any, any, 3) The xyz coordinates of the chunk triplets : array of shape (any, 3) 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 dask.bool to a true bool 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, ) out = np.empty((xyz.shape[0], triplets.shape[0]), dtype=np.float32) if opt: if periodic and box is not None: _geometry._angle_mic(xyz, triplets, box.transpose(0, 2, 1).copy(), out, orthogonal) else: _geometry._angle(xyz, triplets, out) else: out = _angle(xyz, triplets, periodic, out).compute() return out
def compute_angles(traj, angle_indices, periodic=True, opt=True): """Compute the bond angles between the supplied triplets of indices in each frame of a trajectory. Parameters ---------- traj : Trajectory An mtraj trajectory. angle_indices : np.ndarray, shape=(num_pairs, 2), dtype=int Each row gives the indices of three atoms which together make an angle. periodic : bool, default=True If `periodic` is True and the trajectory contains unitcell information, we will treat angles that cross periodic images using the minimum image convention. opt : bool, default=True Use an optimized native library to calculate distances. Our optimized SSE angle calculation implementation is 10-20x faster than the (itself optimized) numpy implementation. Returns ------- angles : np.ndarray, shape=[n_frames, n_angles], dtype=float The angles are in radians """ xyz = ensure_type(traj.xyz, dtype=np.float32, ndim=3, name='traj.xyz', shape=(None, None, 3), warn_on_cast=False) triplets = ensure_type(np.asarray(angle_indices), dtype=np.int32, ndim=2, name='angle_indices', shape=(None, 3), warn_on_cast=False) if not np.all(np.logical_and(triplets < traj.n_atoms, triplets >= 0)): raise ValueError('angle_indices must be between 0 and %d' % traj.n_atoms) out = np.zeros((xyz.shape[0], triplets.shape[0]), dtype=np.float32) 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(): _geometry._angle_mic(xyz, triplets, box, out) return out else: _angle(traj, triplets, periodic, out) return out if opt and _geometry._processor_supports_sse41(): _geometry._angle(xyz, triplets, out) else: _angle(traj, triplets, periodic, out) return out