Package Bio :: Package KDTree :: Module KDTree'
[hide private]
[frames] | no frames]

Source Code for Module Bio.KDTree.KDTree'

  1  # Copyright 2004 by Thomas Hamelryck. 
  2  # All rights reserved. 
  3  # This code is part of the Biopython distribution and governed by its 
  4  # license.  Please see the LICENSE file that should have been included 
  5  # as part of this package. 
  6  """KD tree data structure for searching N-dimensional vectors. 
  7   
  8  The KD tree data structure can be used for all kinds of searches that 
  9  involve N-dimensional vectors, e.g.  neighbor searches (find all points 
 10  within a radius of a given point) or finding all point pairs in a set 
 11  that are within a certain radius of each other. See "Computational Geometry: 
 12  Algorithms and Applications" (Mark de Berg, Marc van Kreveld, Mark Overmars, 
 13  Otfried Schwarzkopf). Author: Thomas Hamelryck. 
 14  """ 
 15   
 16  from __future__ import print_function 
 17   
 18  from numpy import sum, sqrt, array 
 19  from numpy.random import random 
 20   
 21  from Bio.KDTree import _CKDTree 
 22   
 23   
24 -def _dist(p, q):
25 diff = p - q 26 return sqrt(sum(diff * diff))
27 28
29 -def _neighbor_test(nr_points, dim, bucket_size, radius):
30 """ Test all fixed radius neighbor search. 31 32 Test all fixed radius neighbor search using the 33 KD tree C module. 34 35 o nr_points - number of points used in test 36 o dim - dimension of coords 37 o bucket_size - nr of points per tree node 38 o radius - radius of search (typically 0.05 or so) 39 """ 40 # KD tree search 41 kdt = _CKDTree.KDTree(dim, bucket_size) 42 coords = random((nr_points, dim)) 43 kdt.set_data(coords) 44 neighbors = kdt.neighbor_search(radius) 45 r = [neighbor.radius for neighbor in neighbors] 46 if r is None: 47 l1 = 0 48 else: 49 l1 = len(r) 50 # now do a slow search to compare results 51 neighbors = kdt.neighbor_simple_search(radius) 52 r = [neighbor.radius for neighbor in neighbors] 53 if r is None: 54 l2 = 0 55 else: 56 l2 = len(r) 57 if l1 == l2: 58 print("Passed.") 59 else: 60 print("Not passed: %i != %i." % (l1, l2))
61 62
63 -def _test(nr_points, dim, bucket_size, radius):
64 """Test neighbor search. 65 66 Test neighbor search using the KD tree C module. 67 68 o nr_points - number of points used in test 69 o dim - dimension of coords 70 o bucket_size - nr of points per tree node 71 o radius - radius of search (typically 0.05 or so) 72 """ 73 # kd tree search 74 kdt = _CKDTree.KDTree(dim, bucket_size) 75 coords = random((nr_points, dim)) 76 center = coords[0] 77 kdt.set_data(coords) 78 kdt.search_center_radius(center, radius) 79 r = kdt.get_indices() 80 if r is None: 81 l1 = 0 82 else: 83 l1 = len(r) 84 l2 = 0 85 # now do a manual search to compare results 86 for i in range(0, nr_points): 87 p = coords[i] 88 if _dist(p, center) <= radius: 89 l2 = l2 + 1 90 if l1 == l2: 91 print("Passed.") 92 else: 93 print("Not passed: %i != %i." % (l1, l2))
94 95
96 -class KDTree(object):
97 """ 98 KD tree implementation (C++, SWIG python wrapper) 99 100 The KD tree data structure can be used for all kinds of searches that 101 involve N-dimensional vectors, e.g. neighbor searches (find all points 102 within a radius of a given point) or finding all point pairs in a set 103 that are within a certain radius of each other. 104 105 Reference: 106 107 Computational Geometry: Algorithms and Applications 108 Second Edition 109 Mark de Berg, Marc van Kreveld, Mark Overmars, Otfried Schwarzkopf 110 published by Springer-Verlag 111 2nd rev. ed. 2000. 112 ISBN: 3-540-65620-0 113 114 The KD tree data structure is described in chapter 5, pg. 99. 115 116 The following article made clear to me that the nodes should 117 contain more than one point (this leads to dramatic speed 118 improvements for the "all fixed radius neighbor search", see 119 below): 120 121 JL Bentley, "Kd trees for semidynamic point sets," in Sixth Annual ACM 122 Symposium on Computational Geometry, vol. 91. San Francisco, 1990 123 124 This KD implementation also performs a "all fixed radius neighbor search", 125 i.e. it can find all point pairs in a set that are within a certain radius 126 of each other. As far as I know the algorithm has not been published. 127 """ 128
129 - def __init__(self, dim, bucket_size=1):
130 self.dim = dim 131 self.kdt = _CKDTree.KDTree(dim, bucket_size) 132 self.built = 0
133 134 # Set data 135
136 - def set_coords(self, coords):
137 """Add the coordinates of the points. 138 139 o coords - two dimensional NumPy array. E.g. if the points 140 have dimensionality D and there are N points, the coords 141 array should be NxD dimensional. 142 """ 143 if coords.min() <= -1e6 or coords.max() >= 1e6: 144 raise Exception("Points should lie between -1e6 and 1e6") 145 if len(coords.shape) != 2 or coords.shape[1] != self.dim: 146 raise Exception("Expected a Nx%i NumPy array" % self.dim) 147 self.kdt.set_data(coords) 148 self.built = 1
149 150 # Fixed radius search for a point 151
152 - def search(self, center, radius):
153 """Search all points within radius of center. 154 155 o center - one dimensional NumPy array. E.g. if the points have 156 dimensionality D, the center array should be D dimensional. 157 o radius - float>0 158 """ 159 if not self.built: 160 raise Exception("No point set specified") 161 if center.shape != (self.dim,): 162 raise Exception("Expected a %i-dimensional NumPy array" 163 % self.dim) 164 self.kdt.search_center_radius(center, radius)
165
166 - def get_radii(self):
167 """Return radii. 168 169 Return the list of distances from center after 170 a neighbor search. 171 """ 172 a = self.kdt.get_radii() 173 if a is None: 174 return [] 175 return a
176
177 - def get_indices(self):
178 """Return the list of indices. 179 180 Return the list of indices after a neighbor search. 181 The indices refer to the original coords NumPy array. The 182 coordinates with these indices were within radius of center. 183 184 For an index pair, the first index<second index. 185 """ 186 a = self.kdt.get_indices() 187 if a is None: 188 return [] 189 return a
190 191 # Fixed radius search for all points 192
193 - def all_search(self, radius):
194 """All fixed neighbor search. 195 196 Search all point pairs that are within radius. 197 198 o radius - float (>0) 199 """ 200 if not self.built: 201 raise Exception("No point set specified") 202 self.neighbors = self.kdt.neighbor_search(radius)
203
204 - def all_get_indices(self):
205 """Return All Fixed Neighbor Search results. 206 207 Return a Nx2 dim NumPy array containing 208 the indices of the point pairs, where N 209 is the number of neighbor pairs. 210 """ 211 a = array([[neighbor.index1, neighbor.index2] for neighbor in self.neighbors]) 212 return a
213
214 - def all_get_radii(self):
215 """Return All Fixed Neighbor Search results. 216 217 Return an N-dim array containing the distances 218 of all the point pairs, where N is the number 219 of neighbor pairs.. 220 """ 221 return [neighbor.radius for neighbor in self.neighbors]
222 223 if __name__ == "__main__": 224 225 nr_points = 100000 226 dim = 3 227 bucket_size = 10 228 query_radius = 10 229 230 coords = (200 * random((nr_points, dim))) 231 232 kdtree = KDTree(dim, bucket_size) 233 234 # enter coords 235 kdtree.set_coords(coords) 236 237 # Find all point pairs within radius 238 239 kdtree.all_search(query_radius) 240 241 # get indices & radii of points 242 243 # indices is a list of tuples. Each tuple contains the 244 # two indices of a point pair within query_radius of 245 # each other. 246 indices = kdtree.all_get_indices() 247 radii = kdtree.all_get_radii() 248 249 print("Found %i point pairs within radius %f." % (len(indices), query_radius)) 250 251 # Do 10 individual queries 252 253 for i in range(0, 10): 254 # pick a random center 255 center = random(dim) 256 257 # search neighbors 258 kdtree.search(center, query_radius) 259 260 # get indices & radii of points 261 indices = kdtree.get_indices() 262 radii = kdtree.get_radii() 263 264 x, y, z = center 265 print("Found %i points in radius %f around center (%.2f, %.2f, %.2f)." % (len(indices), query_radius, x, y, z)) 266