forked from mdtraj/mdtraj
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_neighbors.py
53 lines (42 loc) · 1.8 KB
/
test_neighbors.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
import numpy as np
import mdtraj as md
from mdtraj.testing import eq
random = np.random.RandomState(0)
def compute_neighbors_reference(traj, cutoff, query_indices, haystack_indices=None):
if haystack_indices is None:
haystack_indices = range(traj.n_atoms)
# explicitly enumerate the pairs of query-haystack indices we need to
# check
pairs = np.array([(q, i) for i in haystack_indices for q in query_indices if i != q])
dists = md.compute_distances(traj, pairs)
# some of the haystack might be within cutoff of more than one of the
# query atoms, so we need unique
reference = [np.unique(pairs[dists[i] < cutoff, 1]) for i in range(traj.n_frames)]
return reference
def test_compute_neighbors_1():
n_frames = 2
n_atoms = 20
cutoff = 2
xyz = random.randn(n_frames, n_atoms, 3)
traj = md.Trajectory(xyz=xyz, topology=None)
query_indices = [0, 1]
value = md.compute_neighbors(traj, cutoff, query_indices)
reference = compute_neighbors_reference(traj, cutoff, query_indices)
for i in range(n_frames):
eq(value[i], reference[i])
def test_compute_neighbors_2(get_fn):
traj = md.load(get_fn('4ZUO.pdb'))
query_indices = traj.top.select('residue 1')
cutoff = 1.0
value = md.compute_neighbors(traj, cutoff, query_indices)
reference = compute_neighbors_reference(traj, cutoff, query_indices)
for i in range(traj.n_frames):
eq(value[i], reference[i])
def test_compute_neighbors_3(get_fn):
traj = md.load(get_fn('test_good.nc'), top=get_fn('test.parm7'))
query_indices = traj.top.select('residue 1')
cutoff = 1.0
value = md.compute_neighbors(traj, cutoff, query_indices)
reference = compute_neighbors_reference(traj, cutoff, query_indices)
for i in range(traj.n_frames):
eq(value[i], reference[i])