1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
|
import numpy as np
import scipy.special
from scipy.spatial.distance import euclidean, pdist, squareform
import mdtraj as md
from mdtraj.geometry import compute_drid
# To keep backwards compatibility
try:
from scipy.stats import nanmean
except ImportError:
from numpy import nanmean
def test_drid_1():
n_frames = 1
n_atoms = 20
top = md.Topology()
chain = top.add_chain()
residue = top.add_residue("X", chain)
for i in range(n_atoms):
top.add_atom("X", None, residue)
t = md.Trajectory(
xyz=np.random.RandomState(0).randn(n_frames, n_atoms, 3),
topology=top,
)
# t contains no bonds
got = compute_drid(t).reshape(n_frames, n_atoms, 3)
for i in range(n_atoms):
others = set(range(n_atoms)) - {i}
rd = 1 / np.array([euclidean(t.xyz[0, i], t.xyz[0, e]) for e in others])
mean = np.mean(rd)
second = np.mean((rd - mean) ** 2) ** (0.5)
third = scipy.special.cbrt(np.mean((rd - mean) ** 3))
ref = np.array([mean, second, third])
np.testing.assert_array_almost_equal(got[0, i], ref, decimal=5)
def test_drid_2():
n_frames = 3
n_atoms = 11
n_bonds = 5
top = md.Topology()
chain = top.add_chain()
residue = top.add_residue("X", chain)
for i in range(n_atoms):
top.add_atom("X", None, residue)
random = np.random.RandomState(0)
bonds = random.randint(n_atoms, size=(n_bonds, 2))
for a, b in bonds:
top.add_bond(top.atom(a), top.atom(b))
t = md.Trajectory(xyz=random.randn(n_frames, n_atoms, 3), topology=top)
got = compute_drid(t).reshape(n_frames, n_atoms, 3)
for i in range(n_frames):
with np.errstate(divide="ignore"):
# For catching division by zero errors
recip = 1 / squareform(pdist(t.xyz[i]))
recip[np.diag_indices(n=recip.shape[0])] = np.nan
recip[bonds[:, 0], bonds[:, 1]] = np.nan
recip[bonds[:, 1], bonds[:, 0]] = np.nan
mean = nanmean(recip, axis=0)
second = nanmean((recip - mean) ** 2, axis=0) ** (0.5)
third = scipy.special.cbrt(nanmean((recip - mean) ** 3, axis=0))
np.testing.assert_array_almost_equal(got[i, :, 0], mean, decimal=5)
np.testing.assert_array_almost_equal(got[i, :, 1], second, decimal=5)
np.testing.assert_array_almost_equal(got[i, :, 2], third, decimal=5)
|
