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