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