forked from mdtraj/mdtraj
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtest_rmsd.py
196 lines (149 loc) · 6.62 KB
/
test_rmsd.py
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
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
##############################################################################
# MDTraj: A Python Library for Loading, Saving, and Manipulating
# Molecular Dynamics Trajectories.
# Copyright 2012-2013 Stanford University and the Authors
#
# Authors: Robert McGibbon
# Contributors:
#
# MDTraj is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 2.1
# of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with MDTraj. If not, see <http://www.gnu.org/licenses/>.
##############################################################################
import mdtraj as md
import numpy as np
from mdtraj.geometry.alignment import rmsd_qcp, compute_average_structure
from mdtraj.testing import eq
np.random.seed(52)
def test_trajectory_rmsd(get_fn):
t = md.load(get_fn('traj.h5'))
for parallel in [True, False]:
calculated = md.rmsd(t, t, 0, parallel=parallel)
reference = np.zeros(t.n_frames)
for i in range(t.n_frames):
reference[i] = rmsd_qcp(t.xyz[0], t.xyz[i])
eq(calculated, reference, decimal=3)
def test_precentered_1(get_fn):
# test rmsd against the numpy version, using the same trajectory
# as target and reference
t1 = md.load(get_fn('traj.h5'), stride=10)
t2 = md.load(get_fn('traj.h5'), stride=10)
# don't center t1, and use it without precentered
# explicitly center t2, and use *with* precentered
for parallel in [True, False]:
t2.center_coordinates()
eq(t1.n_frames, t2.n_frames)
for i in range(t1.n_frames):
ref = np.zeros(t1.n_frames)
for j in range(t1.n_frames):
ref[j] = rmsd_qcp(t1.xyz[j], t1.xyz[i])
val1 = md.rmsd(t1, t1, i, parallel=parallel, precentered=False)
val2 = md.rmsd(t2, t2, i, parallel=parallel, precentered=True)
eq(ref, val1, decimal=3)
eq(val1, val2)
def test_precentered_2(get_fn):
# test rmsd against the numpy version, using the difference
# trajectories as target and reference
t1_a = md.load(get_fn('traj.h5'), stride=10)
t2_a = md.load(get_fn('traj.h5'), stride=10)
t1_b = md.load(get_fn('traj.h5'), stride=10)
t2_b = md.load(get_fn('traj.h5'), stride=10)
# don't center t1, and use it without precentered
# explicitly center t2, and use *with* precentered
t2_a.center_coordinates()
t2_b.center_coordinates()
for parallel in [True, False]:
for i in range(t1_b.n_frames):
ref = np.zeros(t1_a.n_frames)
for j in range(t1_a.n_frames):
ref[j] = rmsd_qcp(t1_a.xyz[j], t1_b.xyz[i])
val1 = md.rmsd(t1_a, t1_b, i, parallel=parallel, precentered=False)
val2 = md.rmsd(t2_a, t2_b, i, parallel=parallel, precentered=True)
eq(ref, val1, decimal=3)
eq(val1, val2, decimal=4)
def test_superpose_0(get_fn):
t1 = md.load(get_fn('traj.h5'))
reference_rmsd = md.rmsd(t1, t1, 0)
t1.superpose(t1, 0)
displ_rmsd = np.zeros(t1.n_frames)
for i in range(t1.n_frames):
delta = t1.xyz[i] - t1.xyz[0]
displ_rmsd[i] = (delta ** 2.0).sum(1).mean() ** 0.5
eq(reference_rmsd, displ_rmsd, decimal=5)
def test_superpose_1():
# make one frame far from the origin
reference = md.Trajectory(xyz=np.random.randn(1, 100, 3) + 100, topology=None)
reference_xyz = reference.xyz.copy()
for indices in [None, np.arange(90)]:
# make another trajectory in a similar rotational state
query = md.Trajectory(xyz=reference.xyz + 0.01 * np.random.randn(*reference.xyz.shape), topology=None)
query.superpose(reference, 0, atom_indices=indices)
assert eq(reference.xyz, reference_xyz)
new_centers = np.mean(query.xyz[0], axis=1)
assert 80 < new_centers[0] < 120
assert 80 < new_centers[1] < 120
assert 80 < new_centers[2] < 120
def test_superpose_2():
t1 = md.Trajectory(xyz=np.random.randn(1, 100, 3) + 100, topology=None)
t2 = md.Trajectory(xyz=np.random.randn(1, 100, 3) + 100, topology=None)
t2_copy = t2.xyz.copy()
t1.superpose(t2)
t1.superpose(t2, atom_indices=[1, 2, 3, 4, 5, 6, 7])
# make sure that superposing doesn't alter the reference traj
eq(t2.xyz, t2_copy)
def test_superpose_refinds():
# make one frame far from the origin
normal = np.random.randn(1, 100, 3)
normal_xyz = normal.copy()
flipped = np.zeros_like(normal)
flipped[:, :50, :] = normal[:, 50:, :]
flipped[:, 50:, :] = normal[:, :50, :]
flipped_xyz = flipped.copy()
normal = md.Trajectory(xyz=normal, topology=None)
flipped = md.Trajectory(xyz=flipped, topology=None)
normal.superpose(flipped, atom_indices=np.arange(0, 50), ref_atom_indices=np.arange(50, 100))
eq(normal.xyz, normal_xyz)
flipped.superpose(normal, atom_indices=np.arange(50, 100), ref_atom_indices=np.arange(0, 50))
eq(flipped.xyz, flipped_xyz)
normal.superpose(flipped)
assert not np.allclose(normal.xyz, normal_xyz)
def test_rmsd_atom_indices(get_fn):
native = md.load(get_fn('native.pdb'))
t1 = md.load(get_fn('traj.h5'))
atom_indices = np.arange(10)
dist1 = md.rmsd(t1, native, atom_indices=atom_indices)
t2 = md.load(get_fn('traj.h5'))
t2.restrict_atoms(atom_indices)
native.restrict_atoms(atom_indices)
dist2 = md.rmsd(t2, native)
eq(dist1, dist2)
def test_rmsd_ref_ainds(get_fn):
native = md.load(get_fn('native.pdb'))
t1 = md.load(get_fn('traj.h5'))
atom_indices = np.arange(10)
dist1 = md.rmsd(t1, native, atom_indices=atom_indices,
ref_atom_indices=atom_indices)
bad_atom_indices = np.arange(10, 20)
t2 = md.load(get_fn('traj.h5'))
dist2 = md.rmsd(t2, native, atom_indices=atom_indices,
ref_atom_indices=bad_atom_indices)
assert np.all(dist2 > dist1)
def test_average_structure(get_fn):
traj = md.load(get_fn('frame0.dcd'), top=get_fn('frame0.pdb'))
average = compute_average_structure(traj.xyz)
# The mean RMSD to the average structure should be less than to any individual frame.
sum1 = 0
sum2 = 0
for i in range(traj.n_frames):
sum1 += rmsd_qcp(traj.xyz[0], traj.xyz[i])
sum2 += rmsd_qcp(average, traj.xyz[i])
assert sum2 < sum1