def _angle(traj, angle_indices, periodic, out): ix01 = angle_indices[:, [1, 0]] ix21 = angle_indices[:, [1, 2]] u_prime = distance.compute_displacements(traj, ix01, periodic=periodic, opt=False) v_prime = distance.compute_displacements(traj, ix21, periodic=periodic, opt=False) u_norm = np.sqrt((u_prime**2).sum(-1)) v_norm = np.sqrt((v_prime**2).sum(-1)) # adding a new axis makes sure that broasting rules kick in on the third # dimension u = u_prime / (u_norm[..., np.newaxis]) v = v_prime / (v_norm[..., np.newaxis]) return np.arccos((u * v).sum(-1), out=out)
def test_2p(): a = compute_distances(ptraj, pairs, periodic=True, opt=False) b = compute_displacements(ptraj, pairs, periodic=True, opt=False) assert a.shape == (len(ptraj), len(pairs)) assert b.shape == (len(ptraj), len(pairs), 3), str(b.shape) b = np.sqrt(np.sum(np.square(b), axis=2)) eq(a, b, decimal=5)
def test_2p(): a = compute_distances(ptraj, pairs, periodic=True, opt=False) b = compute_displacements(ptraj, pairs, periodic=True, opt=False) assert a.shape == (len(ptraj), len(pairs)) assert b.shape == (len(ptraj), len(pairs), 3), str(b.shape) b = np.sqrt(np.sum(np.square(b), axis=2)) eq(a, b, decimal=5)
def _dihedral(traj, indices, periodic, out=None): """Compute the dihedral angles of traj for the atom indices in indices. Parameters ---------- xyz : np.ndarray, shape=(num_frames, num_atoms, 3), dtype=float The XYZ coordinates of a trajectory indices : np.ndarray, shape=(num_dihedrals, 4), dtype=int Atom indices to compute dihedrals. periodic : bool, default=True If `periodic` is True and the trajectory contains unitcell information, we will treat dihedrals that cross periodic images using the minimum image convention. Returns ------- dih : np.ndarray, shape=(num_dihedrals), dtype=float dih[i,j] gives the dihedral angle at traj[i] correponding to indices[j]. """ ix10 = indices[:, [0, 1]] ix21 = indices[:, [1, 2]] ix32 = indices[:, [2, 3]] b1 = distance.compute_displacements(traj, ix10, periodic=periodic, opt=False) b2 = distance.compute_displacements(traj, ix21, periodic=periodic, opt=False) b3 = distance.compute_displacements(traj, ix32, periodic=periodic, opt=False) c1 = np.cross(b2, b3) c2 = np.cross(b1, b2) p1 = (b1 * c1).sum(-1) p1 *= (b2 * b2).sum(-1)**0.5 p2 = (c1 * c2).sum(-1) return np.arctan2(p1, p2, out)
def _dihedral(traj, indices, periodic, out=None): """Compute the dihedral angles of traj for the atom indices in indices. Parameters ---------- xyz : np.ndarray, shape=(num_frames, num_atoms, 3), dtype=float The XYZ coordinates of a trajectory indices : np.ndarray, shape=(num_dihedrals, 4), dtype=int Atom indices to compute dihedrals. periodic : bool, default=True If `periodic` is True and the trajectory contains unitcell information, we will treat dihedrals that cross periodic images using the minimum image convention. Returns ------- dih : np.ndarray, shape=(num_dihedrals), dtype=float dih[i,j] gives the dihedral angle at traj[i] correponding to indices[j]. """ ix10 = indices[:, [0, 1]] ix21 = indices[:, [1, 2]] ix32 = indices[:, [2, 3]] b1 = distance.compute_displacements(traj, ix10, periodic=periodic, opt=False) b2 = distance.compute_displacements(traj, ix21, periodic=periodic, opt=False) b3 = distance.compute_displacements(traj, ix32, periodic=periodic, opt=False) c1 = np.cross(b2, b3) c2 = np.cross(b1, b2) p1 = (b1 * c1).sum(-1) p1 *= (b2 * b2).sum(-1) ** 0.5 p2 = (c1 * c2).sum(-1) return np.arctan2(p1, p2, out)
def test_3p(): a = compute_distances(ptraj, pairs, periodic=True, opt=True) b = compute_displacements(ptraj, pairs, periodic=True, opt=True) eq(a, np.sqrt(np.sum(np.square(b), axis=2)))
def test_1p(): a = compute_displacements(ptraj, pairs, periodic=True, opt=True) b = compute_displacements(ptraj, pairs, periodic=True, opt=False) eq(a, b, decimal=3)
def test_1(): a = compute_displacements(ptraj, pairs, periodic=False, opt=True) b = compute_displacements(ptraj, pairs, periodic=False, opt=False) eq(a, b)
def test_3p(): a = compute_distances(ptraj, pairs, periodic=True, opt=True) b = compute_displacements(ptraj, pairs, periodic=True, opt=True) eq(a, np.sqrt(np.sum(np.square(b), axis=2)))
def test_1p(): a = compute_displacements(ptraj, pairs, periodic=True, opt=True) b = compute_displacements(ptraj, pairs, periodic=True, opt=False) eq(a, b, decimal=3)
def test_1(): a = compute_displacements(ptraj, pairs, periodic=False, opt=True) b = compute_displacements(ptraj, pairs, periodic=False, opt=False) eq(a, b)