1
2
3
4
5
6 from Numeric import array
7
8 from Bio.KDTree import *
9 from PDBExceptions import PDBException
10 from Selection import unfold_entities, get_unique_parents, entity_levels, uniqueify
11
12
13 __doc__="Fast atom neighbor lookup using a KD tree (implemented in C++)."
14
15
16
18 """
19 This class can be used for two related purposes:
20
21 1. To find all atoms/residues/chains/models/structures within radius
22 of a given query position.
23
24 2. To find all atoms/residues/chains/models/structures that are within
25 a fixed radius of each other.
26
27 NeighborSearch makes use of the Bio.KDTree C++ module, so it's fast.
28 """
29 - def __init__(self, atom_list, bucket_size=10):
30 """
31 o atom_list - list of atoms. This list is used in the queries.
32 It can contain atoms from different structures.
33 o bucket_size - bucket size of KD tree. You can play around
34 with this to optimize speed if you feel like it.
35 """
36 self.atom_list=atom_list
37
38 coord_list=map(lambda a: a.get_coord(), atom_list)
39
40 self.coords=array(coord_list).astype("f")
41 assert(bucket_size>1)
42 assert(self.coords.shape[1]==3)
43 assert(self.coords.typecode()=="f")
44 self.kdt=KDTree(3, bucket_size)
45 self.kdt.set_coords(self.coords)
46
47
48
50
51
52
53
54
55 parent_pair_list=[]
56 for (e1, e2) in pair_list:
57 p1=e1.get_parent()
58 p2=e2.get_parent()
59 if p1==p2:
60 continue
61 elif p1<p2:
62 parent_pair_list.append((p1, p2))
63 else:
64 parent_pair_list.append((p2, p1))
65 return uniqueify(parent_pair_list)
66
67
68
69 - def search(self, center, radius, level="A"):
70 """Neighbor search.
71
72 Return all atoms/residues/chains/models/structures
73 that have at least one atom within radius of center.
74 What entitity level is returned (e.g. atoms or residues)
75 is determined by level (A=atoms, R=residues, C=chains,
76 M=models, S=structures).
77
78 o center - Numeric array
79 o radius - float
80 o level - char (A, R, C, M, S)
81 """
82 if not level in entity_levels:
83 raise PDBException, "%s: Unknown level" % level
84 self.kdt.search(center, radius)
85 indices=self.kdt.get_indices()
86 n_atom_list=[]
87 atom_list=self.atom_list
88 for i in indices:
89 a=atom_list[i]
90 n_atom_list.append(a)
91 if level=="A":
92 return n_atom_list
93 else:
94 return unfold_entities(n_atom_list, level)
95
97 """All neighbor search.
98
99 Search all entities that have atoms pairs within
100 radius.
101
102 o radius - float
103 o level - char (A, R, C, M, S)
104 """
105 if not level in entity_levels:
106 raise PDBException, "%s: Unknown level" % level
107 self.kdt.all_search(radius)
108 indices=self.kdt.all_get_indices()
109 atom_list=self.atom_list
110 atom_pair_list=[]
111 for i1, i2 in indices:
112 a1=atom_list[i1]
113 a2=atom_list[i2]
114 atom_pair_list.append((a1, a2))
115 if level=="A":
116
117 return atom_pair_list
118 next_level_pair_list=atom_pair_list
119 for l in ["R", "C", "M", "S"]:
120 next_level_pair_list=self._get_unique_parent_pairs(next_level_pair_list)
121 if level==l:
122 return next_level_pair_list
123
124 if __name__=="__main__":
125
126 from RandomArray import *
127
130 self.coord=(100*random(3))
131
134
135 for i in range(0, 20):
136 al=[]
137 for i in range(0, 100):
138 al.append(Atom())
139
140 ns=NeighborSearch(al)
141
142 print "Found ", len(ns.search_all(5.0))
143