kdtree.py 37 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987
  1. # Copyright Anne M. Archibald 2008
  2. # Released under the scipy license
  3. from __future__ import division, print_function, absolute_import
  4. import sys
  5. import numpy as np
  6. from heapq import heappush, heappop
  7. import scipy.sparse
  8. __all__ = ['minkowski_distance_p', 'minkowski_distance',
  9. 'distance_matrix',
  10. 'Rectangle', 'KDTree']
  11. def minkowski_distance_p(x, y, p=2):
  12. """
  13. Compute the p-th power of the L**p distance between two arrays.
  14. For efficiency, this function computes the L**p distance but does
  15. not extract the pth root. If `p` is 1 or infinity, this is equal to
  16. the actual L**p distance.
  17. Parameters
  18. ----------
  19. x : (M, K) array_like
  20. Input array.
  21. y : (N, K) array_like
  22. Input array.
  23. p : float, 1 <= p <= infinity
  24. Which Minkowski p-norm to use.
  25. Examples
  26. --------
  27. >>> from scipy.spatial import minkowski_distance_p
  28. >>> minkowski_distance_p([[0,0],[0,0]], [[1,1],[0,1]])
  29. array([2, 1])
  30. """
  31. x = np.asarray(x)
  32. y = np.asarray(y)
  33. if p == np.inf:
  34. return np.amax(np.abs(y-x), axis=-1)
  35. elif p == 1:
  36. return np.sum(np.abs(y-x), axis=-1)
  37. else:
  38. return np.sum(np.abs(y-x)**p, axis=-1)
  39. def minkowski_distance(x, y, p=2):
  40. """
  41. Compute the L**p distance between two arrays.
  42. Parameters
  43. ----------
  44. x : (M, K) array_like
  45. Input array.
  46. y : (N, K) array_like
  47. Input array.
  48. p : float, 1 <= p <= infinity
  49. Which Minkowski p-norm to use.
  50. Examples
  51. --------
  52. >>> from scipy.spatial import minkowski_distance
  53. >>> minkowski_distance([[0,0],[0,0]], [[1,1],[0,1]])
  54. array([ 1.41421356, 1. ])
  55. """
  56. x = np.asarray(x)
  57. y = np.asarray(y)
  58. if p == np.inf or p == 1:
  59. return minkowski_distance_p(x, y, p)
  60. else:
  61. return minkowski_distance_p(x, y, p)**(1./p)
  62. class Rectangle(object):
  63. """Hyperrectangle class.
  64. Represents a Cartesian product of intervals.
  65. """
  66. def __init__(self, maxes, mins):
  67. """Construct a hyperrectangle."""
  68. self.maxes = np.maximum(maxes,mins).astype(float)
  69. self.mins = np.minimum(maxes,mins).astype(float)
  70. self.m, = self.maxes.shape
  71. def __repr__(self):
  72. return "<Rectangle %s>" % list(zip(self.mins, self.maxes))
  73. def volume(self):
  74. """Total volume."""
  75. return np.prod(self.maxes-self.mins)
  76. def split(self, d, split):
  77. """
  78. Produce two hyperrectangles by splitting.
  79. In general, if you need to compute maximum and minimum
  80. distances to the children, it can be done more efficiently
  81. by updating the maximum and minimum distances to the parent.
  82. Parameters
  83. ----------
  84. d : int
  85. Axis to split hyperrectangle along.
  86. split : float
  87. Position along axis `d` to split at.
  88. """
  89. mid = np.copy(self.maxes)
  90. mid[d] = split
  91. less = Rectangle(self.mins, mid)
  92. mid = np.copy(self.mins)
  93. mid[d] = split
  94. greater = Rectangle(mid, self.maxes)
  95. return less, greater
  96. def min_distance_point(self, x, p=2.):
  97. """
  98. Return the minimum distance between input and points in the hyperrectangle.
  99. Parameters
  100. ----------
  101. x : array_like
  102. Input.
  103. p : float, optional
  104. Input.
  105. """
  106. return minkowski_distance(0, np.maximum(0,np.maximum(self.mins-x,x-self.maxes)),p)
  107. def max_distance_point(self, x, p=2.):
  108. """
  109. Return the maximum distance between input and points in the hyperrectangle.
  110. Parameters
  111. ----------
  112. x : array_like
  113. Input array.
  114. p : float, optional
  115. Input.
  116. """
  117. return minkowski_distance(0, np.maximum(self.maxes-x,x-self.mins),p)
  118. def min_distance_rectangle(self, other, p=2.):
  119. """
  120. Compute the minimum distance between points in the two hyperrectangles.
  121. Parameters
  122. ----------
  123. other : hyperrectangle
  124. Input.
  125. p : float
  126. Input.
  127. """
  128. return minkowski_distance(0, np.maximum(0,np.maximum(self.mins-other.maxes,other.mins-self.maxes)),p)
  129. def max_distance_rectangle(self, other, p=2.):
  130. """
  131. Compute the maximum distance between points in the two hyperrectangles.
  132. Parameters
  133. ----------
  134. other : hyperrectangle
  135. Input.
  136. p : float, optional
  137. Input.
  138. """
  139. return minkowski_distance(0, np.maximum(self.maxes-other.mins,other.maxes-self.mins),p)
  140. class KDTree(object):
  141. """
  142. kd-tree for quick nearest-neighbor lookup
  143. This class provides an index into a set of k-dimensional points which
  144. can be used to rapidly look up the nearest neighbors of any point.
  145. Parameters
  146. ----------
  147. data : (N,K) array_like
  148. The data points to be indexed. This array is not copied, and
  149. so modifying this data will result in bogus results.
  150. leafsize : int, optional
  151. The number of points at which the algorithm switches over to
  152. brute-force. Has to be positive.
  153. Raises
  154. ------
  155. RuntimeError
  156. The maximum recursion limit can be exceeded for large data
  157. sets. If this happens, either increase the value for the `leafsize`
  158. parameter or increase the recursion limit by::
  159. >>> import sys
  160. >>> sys.setrecursionlimit(10000)
  161. See Also
  162. --------
  163. cKDTree : Implementation of `KDTree` in Cython
  164. Notes
  165. -----
  166. The algorithm used is described in Maneewongvatana and Mount 1999.
  167. The general idea is that the kd-tree is a binary tree, each of whose
  168. nodes represents an axis-aligned hyperrectangle. Each node specifies
  169. an axis and splits the set of points based on whether their coordinate
  170. along that axis is greater than or less than a particular value.
  171. During construction, the axis and splitting point are chosen by the
  172. "sliding midpoint" rule, which ensures that the cells do not all
  173. become long and thin.
  174. The tree can be queried for the r closest neighbors of any given point
  175. (optionally returning only those within some maximum distance of the
  176. point). It can also be queried, with a substantial gain in efficiency,
  177. for the r approximate closest neighbors.
  178. For large dimensions (20 is already large) do not expect this to run
  179. significantly faster than brute force. High-dimensional nearest-neighbor
  180. queries are a substantial open problem in computer science.
  181. The tree also supports all-neighbors queries, both with arrays of points
  182. and with other kd-trees. These do use a reasonably efficient algorithm,
  183. but the kd-tree is not necessarily the best data structure for this
  184. sort of calculation.
  185. """
  186. def __init__(self, data, leafsize=10):
  187. self.data = np.asarray(data)
  188. self.n, self.m = np.shape(self.data)
  189. self.leafsize = int(leafsize)
  190. if self.leafsize < 1:
  191. raise ValueError("leafsize must be at least 1")
  192. self.maxes = np.amax(self.data,axis=0)
  193. self.mins = np.amin(self.data,axis=0)
  194. self.tree = self.__build(np.arange(self.n), self.maxes, self.mins)
  195. class node(object):
  196. if sys.version_info[0] >= 3:
  197. def __lt__(self, other):
  198. return id(self) < id(other)
  199. def __gt__(self, other):
  200. return id(self) > id(other)
  201. def __le__(self, other):
  202. return id(self) <= id(other)
  203. def __ge__(self, other):
  204. return id(self) >= id(other)
  205. def __eq__(self, other):
  206. return id(self) == id(other)
  207. class leafnode(node):
  208. def __init__(self, idx):
  209. self.idx = idx
  210. self.children = len(idx)
  211. class innernode(node):
  212. def __init__(self, split_dim, split, less, greater):
  213. self.split_dim = split_dim
  214. self.split = split
  215. self.less = less
  216. self.greater = greater
  217. self.children = less.children+greater.children
  218. def __build(self, idx, maxes, mins):
  219. if len(idx) <= self.leafsize:
  220. return KDTree.leafnode(idx)
  221. else:
  222. data = self.data[idx]
  223. # maxes = np.amax(data,axis=0)
  224. # mins = np.amin(data,axis=0)
  225. d = np.argmax(maxes-mins)
  226. maxval = maxes[d]
  227. minval = mins[d]
  228. if maxval == minval:
  229. # all points are identical; warn user?
  230. return KDTree.leafnode(idx)
  231. data = data[:,d]
  232. # sliding midpoint rule; see Maneewongvatana and Mount 1999
  233. # for arguments that this is a good idea.
  234. split = (maxval+minval)/2
  235. less_idx = np.nonzero(data <= split)[0]
  236. greater_idx = np.nonzero(data > split)[0]
  237. if len(less_idx) == 0:
  238. split = np.amin(data)
  239. less_idx = np.nonzero(data <= split)[0]
  240. greater_idx = np.nonzero(data > split)[0]
  241. if len(greater_idx) == 0:
  242. split = np.amax(data)
  243. less_idx = np.nonzero(data < split)[0]
  244. greater_idx = np.nonzero(data >= split)[0]
  245. if len(less_idx) == 0:
  246. # _still_ zero? all must have the same value
  247. if not np.all(data == data[0]):
  248. raise ValueError("Troublesome data array: %s" % data)
  249. split = data[0]
  250. less_idx = np.arange(len(data)-1)
  251. greater_idx = np.array([len(data)-1])
  252. lessmaxes = np.copy(maxes)
  253. lessmaxes[d] = split
  254. greatermins = np.copy(mins)
  255. greatermins[d] = split
  256. return KDTree.innernode(d, split,
  257. self.__build(idx[less_idx],lessmaxes,mins),
  258. self.__build(idx[greater_idx],maxes,greatermins))
  259. def __query(self, x, k=1, eps=0, p=2, distance_upper_bound=np.inf):
  260. side_distances = np.maximum(0,np.maximum(x-self.maxes,self.mins-x))
  261. if p != np.inf:
  262. side_distances **= p
  263. min_distance = np.sum(side_distances)
  264. else:
  265. min_distance = np.amax(side_distances)
  266. # priority queue for chasing nodes
  267. # entries are:
  268. # minimum distance between the cell and the target
  269. # distances between the nearest side of the cell and the target
  270. # the head node of the cell
  271. q = [(min_distance,
  272. tuple(side_distances),
  273. self.tree)]
  274. # priority queue for the nearest neighbors
  275. # furthest known neighbor first
  276. # entries are (-distance**p, i)
  277. neighbors = []
  278. if eps == 0:
  279. epsfac = 1
  280. elif p == np.inf:
  281. epsfac = 1/(1+eps)
  282. else:
  283. epsfac = 1/(1+eps)**p
  284. if p != np.inf and distance_upper_bound != np.inf:
  285. distance_upper_bound = distance_upper_bound**p
  286. while q:
  287. min_distance, side_distances, node = heappop(q)
  288. if isinstance(node, KDTree.leafnode):
  289. # brute-force
  290. data = self.data[node.idx]
  291. ds = minkowski_distance_p(data,x[np.newaxis,:],p)
  292. for i in range(len(ds)):
  293. if ds[i] < distance_upper_bound:
  294. if len(neighbors) == k:
  295. heappop(neighbors)
  296. heappush(neighbors, (-ds[i], node.idx[i]))
  297. if len(neighbors) == k:
  298. distance_upper_bound = -neighbors[0][0]
  299. else:
  300. # we don't push cells that are too far onto the queue at all,
  301. # but since the distance_upper_bound decreases, we might get
  302. # here even if the cell's too far
  303. if min_distance > distance_upper_bound*epsfac:
  304. # since this is the nearest cell, we're done, bail out
  305. break
  306. # compute minimum distances to the children and push them on
  307. if x[node.split_dim] < node.split:
  308. near, far = node.less, node.greater
  309. else:
  310. near, far = node.greater, node.less
  311. # near child is at the same distance as the current node
  312. heappush(q,(min_distance, side_distances, near))
  313. # far child is further by an amount depending only
  314. # on the split value
  315. sd = list(side_distances)
  316. if p == np.inf:
  317. min_distance = max(min_distance, abs(node.split-x[node.split_dim]))
  318. elif p == 1:
  319. sd[node.split_dim] = np.abs(node.split-x[node.split_dim])
  320. min_distance = min_distance - side_distances[node.split_dim] + sd[node.split_dim]
  321. else:
  322. sd[node.split_dim] = np.abs(node.split-x[node.split_dim])**p
  323. min_distance = min_distance - side_distances[node.split_dim] + sd[node.split_dim]
  324. # far child might be too far, if so, don't bother pushing it
  325. if min_distance <= distance_upper_bound*epsfac:
  326. heappush(q,(min_distance, tuple(sd), far))
  327. if p == np.inf:
  328. return sorted([(-d,i) for (d,i) in neighbors])
  329. else:
  330. return sorted([((-d)**(1./p),i) for (d,i) in neighbors])
  331. def query(self, x, k=1, eps=0, p=2, distance_upper_bound=np.inf):
  332. """
  333. Query the kd-tree for nearest neighbors
  334. Parameters
  335. ----------
  336. x : array_like, last dimension self.m
  337. An array of points to query.
  338. k : int, optional
  339. The number of nearest neighbors to return.
  340. eps : nonnegative float, optional
  341. Return approximate nearest neighbors; the kth returned value
  342. is guaranteed to be no further than (1+eps) times the
  343. distance to the real kth nearest neighbor.
  344. p : float, 1<=p<=infinity, optional
  345. Which Minkowski p-norm to use.
  346. 1 is the sum-of-absolute-values "Manhattan" distance
  347. 2 is the usual Euclidean distance
  348. infinity is the maximum-coordinate-difference distance
  349. distance_upper_bound : nonnegative float, optional
  350. Return only neighbors within this distance. This is used to prune
  351. tree searches, so if you are doing a series of nearest-neighbor
  352. queries, it may help to supply the distance to the nearest neighbor
  353. of the most recent point.
  354. Returns
  355. -------
  356. d : float or array of floats
  357. The distances to the nearest neighbors.
  358. If x has shape tuple+(self.m,), then d has shape tuple if
  359. k is one, or tuple+(k,) if k is larger than one. Missing
  360. neighbors (e.g. when k > n or distance_upper_bound is
  361. given) are indicated with infinite distances. If k is None,
  362. then d is an object array of shape tuple, containing lists
  363. of distances. In either case the hits are sorted by distance
  364. (nearest first).
  365. i : integer or array of integers
  366. The locations of the neighbors in self.data. i is the same
  367. shape as d.
  368. Examples
  369. --------
  370. >>> from scipy import spatial
  371. >>> x, y = np.mgrid[0:5, 2:8]
  372. >>> tree = spatial.KDTree(list(zip(x.ravel(), y.ravel())))
  373. >>> tree.data
  374. array([[0, 2],
  375. [0, 3],
  376. [0, 4],
  377. [0, 5],
  378. [0, 6],
  379. [0, 7],
  380. [1, 2],
  381. [1, 3],
  382. [1, 4],
  383. [1, 5],
  384. [1, 6],
  385. [1, 7],
  386. [2, 2],
  387. [2, 3],
  388. [2, 4],
  389. [2, 5],
  390. [2, 6],
  391. [2, 7],
  392. [3, 2],
  393. [3, 3],
  394. [3, 4],
  395. [3, 5],
  396. [3, 6],
  397. [3, 7],
  398. [4, 2],
  399. [4, 3],
  400. [4, 4],
  401. [4, 5],
  402. [4, 6],
  403. [4, 7]])
  404. >>> pts = np.array([[0, 0], [2.1, 2.9]])
  405. >>> tree.query(pts)
  406. (array([ 2. , 0.14142136]), array([ 0, 13]))
  407. >>> tree.query(pts[0])
  408. (2.0, 0)
  409. """
  410. x = np.asarray(x)
  411. if np.shape(x)[-1] != self.m:
  412. raise ValueError("x must consist of vectors of length %d but has shape %s" % (self.m, np.shape(x)))
  413. if p < 1:
  414. raise ValueError("Only p-norms with 1<=p<=infinity permitted")
  415. retshape = np.shape(x)[:-1]
  416. if retshape != ():
  417. if k is None:
  418. dd = np.empty(retshape,dtype=object)
  419. ii = np.empty(retshape,dtype=object)
  420. elif k > 1:
  421. dd = np.empty(retshape+(k,),dtype=float)
  422. dd.fill(np.inf)
  423. ii = np.empty(retshape+(k,),dtype=int)
  424. ii.fill(self.n)
  425. elif k == 1:
  426. dd = np.empty(retshape,dtype=float)
  427. dd.fill(np.inf)
  428. ii = np.empty(retshape,dtype=int)
  429. ii.fill(self.n)
  430. else:
  431. raise ValueError("Requested %s nearest neighbors; acceptable numbers are integers greater than or equal to one, or None")
  432. for c in np.ndindex(retshape):
  433. hits = self.__query(x[c], k=k, eps=eps, p=p, distance_upper_bound=distance_upper_bound)
  434. if k is None:
  435. dd[c] = [d for (d,i) in hits]
  436. ii[c] = [i for (d,i) in hits]
  437. elif k > 1:
  438. for j in range(len(hits)):
  439. dd[c+(j,)], ii[c+(j,)] = hits[j]
  440. elif k == 1:
  441. if len(hits) > 0:
  442. dd[c], ii[c] = hits[0]
  443. else:
  444. dd[c] = np.inf
  445. ii[c] = self.n
  446. return dd, ii
  447. else:
  448. hits = self.__query(x, k=k, eps=eps, p=p, distance_upper_bound=distance_upper_bound)
  449. if k is None:
  450. return [d for (d,i) in hits], [i for (d,i) in hits]
  451. elif k == 1:
  452. if len(hits) > 0:
  453. return hits[0]
  454. else:
  455. return np.inf, self.n
  456. elif k > 1:
  457. dd = np.empty(k,dtype=float)
  458. dd.fill(np.inf)
  459. ii = np.empty(k,dtype=int)
  460. ii.fill(self.n)
  461. for j in range(len(hits)):
  462. dd[j], ii[j] = hits[j]
  463. return dd, ii
  464. else:
  465. raise ValueError("Requested %s nearest neighbors; acceptable numbers are integers greater than or equal to one, or None")
  466. def __query_ball_point(self, x, r, p=2., eps=0):
  467. R = Rectangle(self.maxes, self.mins)
  468. def traverse_checking(node, rect):
  469. if rect.min_distance_point(x, p) > r / (1. + eps):
  470. return []
  471. elif rect.max_distance_point(x, p) < r * (1. + eps):
  472. return traverse_no_checking(node)
  473. elif isinstance(node, KDTree.leafnode):
  474. d = self.data[node.idx]
  475. return node.idx[minkowski_distance(d, x, p) <= r].tolist()
  476. else:
  477. less, greater = rect.split(node.split_dim, node.split)
  478. return traverse_checking(node.less, less) + \
  479. traverse_checking(node.greater, greater)
  480. def traverse_no_checking(node):
  481. if isinstance(node, KDTree.leafnode):
  482. return node.idx.tolist()
  483. else:
  484. return traverse_no_checking(node.less) + \
  485. traverse_no_checking(node.greater)
  486. return traverse_checking(self.tree, R)
  487. def query_ball_point(self, x, r, p=2., eps=0):
  488. """Find all points within distance r of point(s) x.
  489. Parameters
  490. ----------
  491. x : array_like, shape tuple + (self.m,)
  492. The point or points to search for neighbors of.
  493. r : positive float
  494. The radius of points to return.
  495. p : float, optional
  496. Which Minkowski p-norm to use. Should be in the range [1, inf].
  497. eps : nonnegative float, optional
  498. Approximate search. Branches of the tree are not explored if their
  499. nearest points are further than ``r / (1 + eps)``, and branches are
  500. added in bulk if their furthest points are nearer than
  501. ``r * (1 + eps)``.
  502. Returns
  503. -------
  504. results : list or array of lists
  505. If `x` is a single point, returns a list of the indices of the
  506. neighbors of `x`. If `x` is an array of points, returns an object
  507. array of shape tuple containing lists of neighbors.
  508. Notes
  509. -----
  510. If you have many points whose neighbors you want to find, you may save
  511. substantial amounts of time by putting them in a KDTree and using
  512. query_ball_tree.
  513. Examples
  514. --------
  515. >>> from scipy import spatial
  516. >>> x, y = np.mgrid[0:5, 0:5]
  517. >>> points = np.c_[x.ravel(), y.ravel()]
  518. >>> tree = spatial.KDTree(points)
  519. >>> tree.query_ball_point([2, 0], 1)
  520. [5, 10, 11, 15]
  521. Query multiple points and plot the results:
  522. >>> import matplotlib.pyplot as plt
  523. >>> points = np.asarray(points)
  524. >>> plt.plot(points[:,0], points[:,1], '.')
  525. >>> for results in tree.query_ball_point(([2, 0], [3, 3]), 1):
  526. ... nearby_points = points[results]
  527. ... plt.plot(nearby_points[:,0], nearby_points[:,1], 'o')
  528. >>> plt.margins(0.1, 0.1)
  529. >>> plt.show()
  530. """
  531. x = np.asarray(x)
  532. if x.shape[-1] != self.m:
  533. raise ValueError("Searching for a %d-dimensional point in a "
  534. "%d-dimensional KDTree" % (x.shape[-1], self.m))
  535. if len(x.shape) == 1:
  536. return self.__query_ball_point(x, r, p, eps)
  537. else:
  538. retshape = x.shape[:-1]
  539. result = np.empty(retshape, dtype=object)
  540. for c in np.ndindex(retshape):
  541. result[c] = self.__query_ball_point(x[c], r, p=p, eps=eps)
  542. return result
  543. def query_ball_tree(self, other, r, p=2., eps=0):
  544. """Find all pairs of points whose distance is at most r
  545. Parameters
  546. ----------
  547. other : KDTree instance
  548. The tree containing points to search against.
  549. r : float
  550. The maximum distance, has to be positive.
  551. p : float, optional
  552. Which Minkowski norm to use. `p` has to meet the condition
  553. ``1 <= p <= infinity``.
  554. eps : float, optional
  555. Approximate search. Branches of the tree are not explored
  556. if their nearest points are further than ``r/(1+eps)``, and
  557. branches are added in bulk if their furthest points are nearer
  558. than ``r * (1+eps)``. `eps` has to be non-negative.
  559. Returns
  560. -------
  561. results : list of lists
  562. For each element ``self.data[i]`` of this tree, ``results[i]`` is a
  563. list of the indices of its neighbors in ``other.data``.
  564. """
  565. results = [[] for i in range(self.n)]
  566. def traverse_checking(node1, rect1, node2, rect2):
  567. if rect1.min_distance_rectangle(rect2, p) > r/(1.+eps):
  568. return
  569. elif rect1.max_distance_rectangle(rect2, p) < r*(1.+eps):
  570. traverse_no_checking(node1, node2)
  571. elif isinstance(node1, KDTree.leafnode):
  572. if isinstance(node2, KDTree.leafnode):
  573. d = other.data[node2.idx]
  574. for i in node1.idx:
  575. results[i] += node2.idx[minkowski_distance(d,self.data[i],p) <= r].tolist()
  576. else:
  577. less, greater = rect2.split(node2.split_dim, node2.split)
  578. traverse_checking(node1,rect1,node2.less,less)
  579. traverse_checking(node1,rect1,node2.greater,greater)
  580. elif isinstance(node2, KDTree.leafnode):
  581. less, greater = rect1.split(node1.split_dim, node1.split)
  582. traverse_checking(node1.less,less,node2,rect2)
  583. traverse_checking(node1.greater,greater,node2,rect2)
  584. else:
  585. less1, greater1 = rect1.split(node1.split_dim, node1.split)
  586. less2, greater2 = rect2.split(node2.split_dim, node2.split)
  587. traverse_checking(node1.less,less1,node2.less,less2)
  588. traverse_checking(node1.less,less1,node2.greater,greater2)
  589. traverse_checking(node1.greater,greater1,node2.less,less2)
  590. traverse_checking(node1.greater,greater1,node2.greater,greater2)
  591. def traverse_no_checking(node1, node2):
  592. if isinstance(node1, KDTree.leafnode):
  593. if isinstance(node2, KDTree.leafnode):
  594. for i in node1.idx:
  595. results[i] += node2.idx.tolist()
  596. else:
  597. traverse_no_checking(node1, node2.less)
  598. traverse_no_checking(node1, node2.greater)
  599. else:
  600. traverse_no_checking(node1.less, node2)
  601. traverse_no_checking(node1.greater, node2)
  602. traverse_checking(self.tree, Rectangle(self.maxes, self.mins),
  603. other.tree, Rectangle(other.maxes, other.mins))
  604. return results
  605. def query_pairs(self, r, p=2., eps=0):
  606. """
  607. Find all pairs of points within a distance.
  608. Parameters
  609. ----------
  610. r : positive float
  611. The maximum distance.
  612. p : float, optional
  613. Which Minkowski norm to use. `p` has to meet the condition
  614. ``1 <= p <= infinity``.
  615. eps : float, optional
  616. Approximate search. Branches of the tree are not explored
  617. if their nearest points are further than ``r/(1+eps)``, and
  618. branches are added in bulk if their furthest points are nearer
  619. than ``r * (1+eps)``. `eps` has to be non-negative.
  620. Returns
  621. -------
  622. results : set
  623. Set of pairs ``(i,j)``, with ``i < j``, for which the corresponding
  624. positions are close.
  625. """
  626. results = set()
  627. def traverse_checking(node1, rect1, node2, rect2):
  628. if rect1.min_distance_rectangle(rect2, p) > r/(1.+eps):
  629. return
  630. elif rect1.max_distance_rectangle(rect2, p) < r*(1.+eps):
  631. traverse_no_checking(node1, node2)
  632. elif isinstance(node1, KDTree.leafnode):
  633. if isinstance(node2, KDTree.leafnode):
  634. # Special care to avoid duplicate pairs
  635. if id(node1) == id(node2):
  636. d = self.data[node2.idx]
  637. for i in node1.idx:
  638. for j in node2.idx[minkowski_distance(d,self.data[i],p) <= r]:
  639. if i < j:
  640. results.add((i,j))
  641. else:
  642. d = self.data[node2.idx]
  643. for i in node1.idx:
  644. for j in node2.idx[minkowski_distance(d,self.data[i],p) <= r]:
  645. if i < j:
  646. results.add((i,j))
  647. elif j < i:
  648. results.add((j,i))
  649. else:
  650. less, greater = rect2.split(node2.split_dim, node2.split)
  651. traverse_checking(node1,rect1,node2.less,less)
  652. traverse_checking(node1,rect1,node2.greater,greater)
  653. elif isinstance(node2, KDTree.leafnode):
  654. less, greater = rect1.split(node1.split_dim, node1.split)
  655. traverse_checking(node1.less,less,node2,rect2)
  656. traverse_checking(node1.greater,greater,node2,rect2)
  657. else:
  658. less1, greater1 = rect1.split(node1.split_dim, node1.split)
  659. less2, greater2 = rect2.split(node2.split_dim, node2.split)
  660. traverse_checking(node1.less,less1,node2.less,less2)
  661. traverse_checking(node1.less,less1,node2.greater,greater2)
  662. # Avoid traversing (node1.less, node2.greater) and
  663. # (node1.greater, node2.less) (it's the same node pair twice
  664. # over, which is the source of the complication in the
  665. # original KDTree.query_pairs)
  666. if id(node1) != id(node2):
  667. traverse_checking(node1.greater,greater1,node2.less,less2)
  668. traverse_checking(node1.greater,greater1,node2.greater,greater2)
  669. def traverse_no_checking(node1, node2):
  670. if isinstance(node1, KDTree.leafnode):
  671. if isinstance(node2, KDTree.leafnode):
  672. # Special care to avoid duplicate pairs
  673. if id(node1) == id(node2):
  674. for i in node1.idx:
  675. for j in node2.idx:
  676. if i < j:
  677. results.add((i,j))
  678. else:
  679. for i in node1.idx:
  680. for j in node2.idx:
  681. if i < j:
  682. results.add((i,j))
  683. elif j < i:
  684. results.add((j,i))
  685. else:
  686. traverse_no_checking(node1, node2.less)
  687. traverse_no_checking(node1, node2.greater)
  688. else:
  689. # Avoid traversing (node1.less, node2.greater) and
  690. # (node1.greater, node2.less) (it's the same node pair twice
  691. # over, which is the source of the complication in the
  692. # original KDTree.query_pairs)
  693. if id(node1) == id(node2):
  694. traverse_no_checking(node1.less, node2.less)
  695. traverse_no_checking(node1.less, node2.greater)
  696. traverse_no_checking(node1.greater, node2.greater)
  697. else:
  698. traverse_no_checking(node1.less, node2)
  699. traverse_no_checking(node1.greater, node2)
  700. traverse_checking(self.tree, Rectangle(self.maxes, self.mins),
  701. self.tree, Rectangle(self.maxes, self.mins))
  702. return results
  703. def count_neighbors(self, other, r, p=2.):
  704. """
  705. Count how many nearby pairs can be formed.
  706. Count the number of pairs (x1,x2) can be formed, with x1 drawn
  707. from self and x2 drawn from `other`, and where
  708. ``distance(x1, x2, p) <= r``.
  709. This is the "two-point correlation" described in Gray and Moore 2000,
  710. "N-body problems in statistical learning", and the code here is based
  711. on their algorithm.
  712. Parameters
  713. ----------
  714. other : KDTree instance
  715. The other tree to draw points from.
  716. r : float or one-dimensional array of floats
  717. The radius to produce a count for. Multiple radii are searched with
  718. a single tree traversal.
  719. p : float, 1<=p<=infinity, optional
  720. Which Minkowski p-norm to use
  721. Returns
  722. -------
  723. result : int or 1-D array of ints
  724. The number of pairs. Note that this is internally stored in a numpy
  725. int, and so may overflow if very large (2e9).
  726. """
  727. def traverse(node1, rect1, node2, rect2, idx):
  728. min_r = rect1.min_distance_rectangle(rect2,p)
  729. max_r = rect1.max_distance_rectangle(rect2,p)
  730. c_greater = r[idx] > max_r
  731. result[idx[c_greater]] += node1.children*node2.children
  732. idx = idx[(min_r <= r[idx]) & (r[idx] <= max_r)]
  733. if len(idx) == 0:
  734. return
  735. if isinstance(node1,KDTree.leafnode):
  736. if isinstance(node2,KDTree.leafnode):
  737. ds = minkowski_distance(self.data[node1.idx][:,np.newaxis,:],
  738. other.data[node2.idx][np.newaxis,:,:],
  739. p).ravel()
  740. ds.sort()
  741. result[idx] += np.searchsorted(ds,r[idx],side='right')
  742. else:
  743. less, greater = rect2.split(node2.split_dim, node2.split)
  744. traverse(node1, rect1, node2.less, less, idx)
  745. traverse(node1, rect1, node2.greater, greater, idx)
  746. else:
  747. if isinstance(node2,KDTree.leafnode):
  748. less, greater = rect1.split(node1.split_dim, node1.split)
  749. traverse(node1.less, less, node2, rect2, idx)
  750. traverse(node1.greater, greater, node2, rect2, idx)
  751. else:
  752. less1, greater1 = rect1.split(node1.split_dim, node1.split)
  753. less2, greater2 = rect2.split(node2.split_dim, node2.split)
  754. traverse(node1.less,less1,node2.less,less2,idx)
  755. traverse(node1.less,less1,node2.greater,greater2,idx)
  756. traverse(node1.greater,greater1,node2.less,less2,idx)
  757. traverse(node1.greater,greater1,node2.greater,greater2,idx)
  758. R1 = Rectangle(self.maxes, self.mins)
  759. R2 = Rectangle(other.maxes, other.mins)
  760. if np.shape(r) == ():
  761. r = np.array([r])
  762. result = np.zeros(1,dtype=int)
  763. traverse(self.tree, R1, other.tree, R2, np.arange(1))
  764. return result[0]
  765. elif len(np.shape(r)) == 1:
  766. r = np.asarray(r)
  767. n, = r.shape
  768. result = np.zeros(n,dtype=int)
  769. traverse(self.tree, R1, other.tree, R2, np.arange(n))
  770. return result
  771. else:
  772. raise ValueError("r must be either a single value or a one-dimensional array of values")
  773. def sparse_distance_matrix(self, other, max_distance, p=2.):
  774. """
  775. Compute a sparse distance matrix
  776. Computes a distance matrix between two KDTrees, leaving as zero
  777. any distance greater than max_distance.
  778. Parameters
  779. ----------
  780. other : KDTree
  781. max_distance : positive float
  782. p : float, optional
  783. Returns
  784. -------
  785. result : dok_matrix
  786. Sparse matrix representing the results in "dictionary of keys" format.
  787. """
  788. result = scipy.sparse.dok_matrix((self.n,other.n))
  789. def traverse(node1, rect1, node2, rect2):
  790. if rect1.min_distance_rectangle(rect2, p) > max_distance:
  791. return
  792. elif isinstance(node1, KDTree.leafnode):
  793. if isinstance(node2, KDTree.leafnode):
  794. for i in node1.idx:
  795. for j in node2.idx:
  796. d = minkowski_distance(self.data[i],other.data[j],p)
  797. if d <= max_distance:
  798. result[i,j] = d
  799. else:
  800. less, greater = rect2.split(node2.split_dim, node2.split)
  801. traverse(node1,rect1,node2.less,less)
  802. traverse(node1,rect1,node2.greater,greater)
  803. elif isinstance(node2, KDTree.leafnode):
  804. less, greater = rect1.split(node1.split_dim, node1.split)
  805. traverse(node1.less,less,node2,rect2)
  806. traverse(node1.greater,greater,node2,rect2)
  807. else:
  808. less1, greater1 = rect1.split(node1.split_dim, node1.split)
  809. less2, greater2 = rect2.split(node2.split_dim, node2.split)
  810. traverse(node1.less,less1,node2.less,less2)
  811. traverse(node1.less,less1,node2.greater,greater2)
  812. traverse(node1.greater,greater1,node2.less,less2)
  813. traverse(node1.greater,greater1,node2.greater,greater2)
  814. traverse(self.tree, Rectangle(self.maxes, self.mins),
  815. other.tree, Rectangle(other.maxes, other.mins))
  816. return result
  817. def distance_matrix(x, y, p=2, threshold=1000000):
  818. """
  819. Compute the distance matrix.
  820. Returns the matrix of all pair-wise distances.
  821. Parameters
  822. ----------
  823. x : (M, K) array_like
  824. Matrix of M vectors in K dimensions.
  825. y : (N, K) array_like
  826. Matrix of N vectors in K dimensions.
  827. p : float, 1 <= p <= infinity
  828. Which Minkowski p-norm to use.
  829. threshold : positive int
  830. If ``M * N * K`` > `threshold`, algorithm uses a Python loop instead
  831. of large temporary arrays.
  832. Returns
  833. -------
  834. result : (M, N) ndarray
  835. Matrix containing the distance from every vector in `x` to every vector
  836. in `y`.
  837. Examples
  838. --------
  839. >>> from scipy.spatial import distance_matrix
  840. >>> distance_matrix([[0,0],[0,1]], [[1,0],[1,1]])
  841. array([[ 1. , 1.41421356],
  842. [ 1.41421356, 1. ]])
  843. """
  844. x = np.asarray(x)
  845. m, k = x.shape
  846. y = np.asarray(y)
  847. n, kk = y.shape
  848. if k != kk:
  849. raise ValueError("x contains %d-dimensional vectors but y contains %d-dimensional vectors" % (k, kk))
  850. if m*n*k <= threshold:
  851. return minkowski_distance(x[:,np.newaxis,:],y[np.newaxis,:,:],p)
  852. else:
  853. result = np.empty((m,n),dtype=float) # FIXME: figure out the best dtype
  854. if m < n:
  855. for i in range(m):
  856. result[i,:] = minkowski_distance(x[i],y,p)
  857. else:
  858. for j in range(n):
  859. result[:,j] = minkowski_distance(x,y[j],p)
  860. return result