Поиск точек в пределах среза расстояния других точек с SciPy

Question

Скажем, у меня есть два набора точек:

>>> points1.shape 
(10000, 3) 
>>> points2.shape 
(1529, 3) 

И я хочу, список индексов points1 в пределах евклидова расстояния cutoff точки в points2. Я мог бы легко сделать это с помощью scipy.spatial.distance.cdist так:

from scipy.spatial.distance import cdist 
import numpy 

indices = numpy.argwhere(cdist(points1, points2).min(axis=0) < cutoff) 

Однако это кажется неэффективным, так как я не нужно знать, как далеко точки друг от друга, просто ли они в пределах расстояния отсечки. Могла ли KDTree помочь в этом?

Answer 1

Вот 3 альтернативы, один с использованием cdist, два с помощью scipy.spatial.cKDTree:

import itertools as IT 
import numpy as np 
import scipy.spatial as spatial 
import scipy.spatial.distance as dist 
np.random.seed(2016) 
points1 = np.random.randint(100, size=(10**5, 3)) 
points2 = np.random.randint(100, size=(1529, 3)) 
cutoff = 5 

def using_cdist(points1, points2, cutoff): 
    indices = np.where(dist.cdist(points1, points2) <= cutoff)[0] 
    indices = np.unique(indices) 
    return indices 

def using_kdtree(points1, points2, cutoff): 
    # build the KDTree using the *smaller* points array 
    tree = spatial.cKDTree(points2) 
    groups = tree.query_ball_point(points1, cutoff) 
    indices = np.unique([i for i, grp in enumerate(groups) if len(grp)]) 
    return indices 

def using_kdtree2(points1, points2, cutoff): 
    # build the KDTree using the *larger* points array 
    tree = spatial.cKDTree(points1) 
    groups = tree.query_ball_point(points2, cutoff) 
    indices = np.unique(IT.chain.from_iterable(groups)) 
    return indices 

cdist_result = using_cdist(points1, points2, cutoff) 
kdtree_result = using_kdtree(points1, points2, cutoff) 
kdtree_result2 = using_kdtree2(points1, points2, cutoff) 
assert np.allclose(cdist_result, kdtree_result) 
assert np.allclose(cdist_result, kdtree_result2) 

Из этих 3-х альтернатив, using_kdtree2 является самым быстрым:

In [80]: %timeit using_kdtree3(points1, points2, cutoff) 
10 loops, best of 3: 92.4 ms per loop 

In [103]: %timeit using_kdtree(points1, points2, cutoff) 
1 loops, best of 3: 938 ms per loop 

In [104]: %timeit using_cdist(points1, points2, cutoff) 
1 loops, best of 3: 1.51 s per loop 

Моя интуиция о том, что будет быть самым быстрым, оказалось совершенно неправильным. I Мысль о создании KDTree с использованием меньшего массива точек будет быстрее всего. Даже если строить KDTree используя больше точек массива является несколько медленнее, вызывая tree.query_ball_point на меньших точек массива гораздо быстрее:

In [68]: %timeit tree = spatial.cKDTree(points2) 
1000 loops, best of 3: 312 µs per loop 

In [69]: %timeit tree = spatial.cKDTree(points1) 
10 loops, best of 3: 45.7 ms per loop 

In [66]: %timeit tree = spatial.cKDTree(points2); groups = tree.query_ball_point(points1, cutoff) 
1 loops, best of 3: 933 ms per loop 

In [67]: %timeit tree = spatial.cKDTree(points1); groups = tree.query_ball_point(points2, cutoff) 
10 loops, best of 3: 89.3 ms per loop 

---

Обратите внимание, что существуют некоторые проблемы, связанные с использованием

def orig(points1, points2, cutoff): 
    return np.argwhere(dist.cdist(points1, points2).min(axis=0) < cutoff) 

Во-первых, позвонив по номеру min(axis=0), вы теряете информацию, если две точки в points1 находятся в пределах cutoff точки в points2. Вы получите только индекс ближайшей точки. Другая проблема заключается в том, что, вызывая min по оси 0, все, что остается, является 1-ой осью, которая связана с points2. Таким образом, orig возвращает индексы в points2, а не points1.

Answer 2

Несколько идей (?):

• Если вам не нужно знать расстояние, вы можете сохранить вычисление квадратного корня, сравнивая квадрат расстояния до квадрата порог (cdist будет вычислять квадратный корень).

• Первый проход, чтобы исключить точки, координаты которых уже превышают пороговое значение, то же самое с y и z, сэкономит некоторые вычисления. Тем более, что «или» ленив в Python (если x уже достаточно далеко, он даже не проверяет y).