有效地查找具有切割的邻居并返回索引

时间:2017-09-18 06:30:42

标签: python numpy indexing nearest-neighbor

我在x,y平面上有很多点,长度大约为10000,每个点(x,y)都有一个内在半径r。这个小数据集只是我整个数据集的一个小角落。我有一个感兴趣的点(x1,y1),我想找到1 {1}附近的附近点,并符合(x1,y1)(x,y)之间的距离小于{{1 }}。我想要返回那些优点的索引,而不是自己的优点。

(x1,y1)

在这个函数中,我可以找到r周围的好点,但不能找到那些好点的索引。但是,我需要ORIGINAL索引,因为ORIGINAL索引用于提取与坐标import numpy as np np.random.seed(2000) x = 20.*np.random.rand(10000) y = 20.*np.random.rand(10000) r = 0.3*np.random.rand(10000) x1 = 10. ### (x1,y1) is an interest point y1 = 12. def index_finder(x,y,r,x1,y1): idx = (abs(x - x1) < 1.) & (abs(y - y1) < 1.) ### This cut will probably cut 90% of the data x_temp = x[idx] ### but if I do like this, then I lose the track of the original index y_temp = y[idx] dis_square = (x_temp - x1)*(x_temp - x1) + (y_temp - y1)*(y_temp - y1) idx1 = dis_square < r*r ### after this cut, there are only a few left x_good = x_temp[idx1] y_good = y_temp[idx1] 关联的其他数据。正如我所提到的,样本数据集只是我整个数据集的一个小角落,我将为我的整个数据集调用上述函数大约1,000,000次,因此上述(x1,y1)函数的效率也是一个考虑因素

对此类任务的任何想法?

3 个答案:

答案 0 :(得分:1)

方法#1

我们可以简单地使用自己的蒙版索引到第一个蒙版,以便从第二个阶段中选择True置位蒙版值,如下所示 -

idx[idx] = idx1

因此,idx将拥有与原始数组xy对应的最终有效屏蔽值/有价值的位置,即 -

x_good = x[idx]
y_good = y[idx]

然后可以使用此掩码索引到问题中提到的其他数组。

方法#2

作为另一种方法,我们可以使用两个条件语句,从而用它们创建两个掩码。最后,将它们与AND-ing组合以获得组合掩码,可以将其索引到最终输出的xy数组中。我们不需要那样获得实际的指数,这样就可以获得另一个好处。

因此,实施 -

X = x-x1
Y = y-y1
mask1 = (np.abs(X) < 1.) & (np.abs(Y) < 1.)
mask2 = X**2 + Y*2 < r**2
comb_mask = mask1 & mask2

x_good = x[comb_mask]
y_good = y[comb_mask]

如果由于某种原因,您仍需要相应的索引,只需执行 -

comb_idx = np.flatnonzero(comb_mask)

如果您对同一x1y1数据集的不同xy对进行这些操作,我建议您使用broadcasting进行矢量化它通过所有x1y1配对数据集,如this post所示。

答案 1 :(得分:1)

numpy.where似乎是为了寻找指数

向量化范数calc + np.where()可能比循环

更快
sq_norm = (x - x1)**2 + (y - y1)**2  # no need to take 10000 sqrt
idcs = np.where(sq_norm < 1.)

len(idcs[0])
Out[193]: 69

np.stack((idcs[0], x[idcs], y[idcs]), axis=1)[:5]
Out[194]: 
array([[  38.        ,    9.47165956,   11.94250173],
       [  39.        ,    9.6966941 ,   11.67505453],
       [ 276.        ,   10.68835317,   12.11589316],
       [ 288.        ,    9.93632584,   11.07624915],
       [ 344.        ,    9.48644057,   12.04911857]])

规范计算也可以包括r数组,第二步?

r_sq_norm = (x[idcs] - x1)**2 + (y[idcs] - y1)**2 - r[idcs]**2
r_idcs = np.where(r_sq_norm < 0.)

idcs[0][r_idcs]
Out[11]: array([1575, 3476, 3709], dtype=int64)

你可能想要在第一个矢量化规范计算中对2步测试和r进行计时吗?

sq_norm = (x - x1)**2 + (y - y1)**2 - r**2
idcs = np.where(sq_norm < 0.)

idcs[0]
Out[13]: array([1575, 3476, 3709], dtype=int64)

答案 2 :(得分:0)

您可以对索引进行掩码,如下所示:

def index_finder(x,y,r,x1,y1):
    idx = np.nonzero((abs(x - x1) < 1.) & (abs(y - y1) < 1.))  #numerical, not boolean
    mask = (x[idx] - x1)*(x[idx] - x1) + (y[idx] - y1)*(y[idx] - y1) < r*r
    idx1 = [i[mask] for i in idx]
    x_good = x_temp[idx1]
    y_good = y_temp[idx1]

现在idx1是您要提取的索引。

一般来说,更快的方法是使用scipy.spatial.KDTree

from scipy.spatial import KDTree

xy = np.stack((x,y))
kdt = KDTree(xy)
kdt.query_ball_point([x1, y1], r)

如果你有很多要点针对同一个数据集进行查询,那么 比顺序调用你的index_finder应用程序更快。

x1y1 = np.stack((x1, y1)) #`x1` and `y1` are arrays of coordinates.
kdt.query_ball_point(x1y1, r)

还错:如果每个点的距离不同,您可以这样做:

def query_variable_ball(kdtree, x, y, r):
    out = []
    for x_, y_, r_ in zip(x, y, r):
        out.append(kdt.query_ball_point([x_, y_], r_)
    return out

xy = np.stack((x,y))
kdt = KDTree(xy)
query_variable_ball(kdt, x1, y1, r)

编辑2:这应该适用于每个点的不同r

from scipy.spatial import KDTree

def index_finder_kd(x, y, r, x1, y1):  # all arrays
    xy = np.stack((x,y), axis = -1)
    x1y1 = np.stack((x1, y1), axis = -1)
    xytree = KDTree(xy)
    d, i = xytree.query(x1y1, k = None, distance_upper_bound = 1.)
    good_idx = np.zeros(x.size, dtype = bool)
    for idx, dist in zip(i, d):
        good_idx[idx] |= r[idx] > dist
    x_good = x[good_idx]
    y_good = y[good_idx]
    return x_good, y_good, np.flatnonzero(good_idx)

由于(x1, y1)需要一段时间来填充,因此只有一个KDTree非常慢。但是如果你有数百万对,那就会快得多。

(我假设你想要所有(x, y)的{​​{1}}数据中的所有优点的联合,如果你想单独使用它们,也可以使用类似的方法,删除{(x1, y1)的元素1}}基于i[j]