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)
