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  from numpy import array, empty 
 18  from Bio.KDTree import _CKDTree 
 19   
 20   
21 -class KDTree(object):
22 """KD tree implementation in C++, SWIG python wrapper. 23 24 The KD tree data structure can be used for all kinds of searches that 25 involve N-dimensional vectors, e.g. neighbor searches (find all points 26 within a radius of a given point) or finding all point pairs in a set 27 that are within a certain radius of each other. 28 29 Reference: 30 31 Computational Geometry: Algorithms and Applications 32 Second Edition 33 Mark de Berg, Marc van Kreveld, Mark Overmars, Otfried Schwarzkopf 34 published by Springer-Verlag 35 2nd rev. ed. 2000. 36 ISBN: 3-540-65620-0 37 38 The KD tree data structure is described in chapter 5, pg. 99. 39 40 The following article made clear to me that the nodes should 41 contain more than one point (this leads to dramatic speed 42 improvements for the "all fixed radius neighbor search", see 43 below): 44 45 JL Bentley, "Kd trees for semidynamic point sets," in Sixth Annual ACM 46 Symposium on Computational Geometry, vol. 91. San Francisco, 1990 47 48 This KD implementation also performs a "all fixed radius neighbor search", 49 i.e. it can find all point pairs in a set that are within a certain radius 50 of each other. As far as I know the algorithm has not been published. 51 """ 52
53 - def __init__(self, dim, bucket_size=1):
54 """Initialize KDTree class.""" 55 self.dim = dim 56 self.kdt = _CKDTree.KDTree(dim, bucket_size) 57 self.built = 0
58 59 # Set data 60
61 - def set_coords(self, coords):
62 """Add the coordinates of the points. 63 64 Arguments: 65 - coords: two dimensional NumPy array. E.g. if the points 66 have dimensionality D and there are N points, the coords 67 array should be NxD dimensional. 68 69 """ 70 if coords.min() <= -1e6 or coords.max() >= 1e6: 71 raise Exception("Points should lie between -1e6 and 1e6") 72 if len(coords.shape) != 2 or coords.shape[1] != self.dim: 73 raise Exception("Expected a Nx%i NumPy array" % self.dim) 74 self.kdt.set_data(coords) 75 self.built = 1
76 77 # Fixed radius search for a point 78
79 - def search(self, center, radius):
80 """Search all points within radius of center. 81 82 Arguments: 83 - center: one dimensional NumPy array. E.g. if the points have 84 dimensionality D, the center array should be D dimensional. 85 - radius: float>0 86 87 """ 88 if not self.built: 89 raise Exception("No point set specified") 90 if center.shape != (self.dim,): 91 raise Exception("Expected a %i-dimensional NumPy array" 92 % self.dim) 93 self.kdt.search_center_radius(center, radius)
94
95 - def get_radii(self):
96 """Return radii. 97 98 Return the list of distances from center after 99 a neighbor search. 100 """ 101 n = self.kdt.get_count() 102 if n == 0: 103 return [] 104 radii = empty(n, int) 105 self.kdt.get_radii(radii) 106 return radii
107
108 - def get_indices(self):
109 """Return the list of indices. 110 111 Return the list of indices after a neighbor search. 112 The indices refer to the original coords NumPy array. The 113 coordinates with these indices were within radius of center. 114 115 For an index pair, the first index<second index. 116 """ 117 n = self.kdt.get_count() 118 if n == 0: 119 return [] 120 indices = empty(n, int) 121 self.kdt.get_indices(indices) 122 return indices
123
124 - def all_search(self, radius):
125 """All fixed neighbor search. 126 127 Search all point pairs that are within radius. 128 129 Arguments: 130 - radius: float (>0) 131 132 """ 133 # Fixed radius search for all points 134 if not self.built: 135 raise Exception("No point set specified") 136 self.neighbors = self.kdt.neighbor_search(radius)
137
138 - def all_get_indices(self):
139 """Return All Fixed Neighbor Search results. 140 141 Return a Nx2 dim NumPy array containing 142 the indices of the point pairs, where N 143 is the number of neighbor pairs. 144 """ 145 a = array([[neighbor.index1, neighbor.index2] for neighbor in self.neighbors]) 146 return a
147
148 - def all_get_radii(self):
149 """Return All Fixed Neighbor Search results. 150 151 Return an N-dim array containing the distances 152 of all the point pairs, where N is the number 153 of neighbor pairs.. 154 """ 155 return [neighbor.radius for neighbor in self.neighbors]
156