用Numpy/Python加速数组查询

Question

我有一个点数组（称为points），由~30000个x、y和z值组成。我还有一个单独的点数组（称为顶点），大约有40000个x、y和z值。后一个数组索引一些边长大小的立方体的左下角。我想知道哪些点位于哪个立方体中，以及每个立方体中有多少个点。我写了一个循环来做这个，它是这样工作的：

for i in xrange(len(vertices)):        
    cube=((vertices[i,0]<= points[:,0]) & 
    (points[:,0]<(vertices[i,0]+size)) & 
    (vertices[i,1]<= points[:,1]) & 
    (points[:,1] < (vertices[i,1]+size)) &
    (vertices[i,2]<= points[:,2]) & 
    (points[:,2] < (vertices[i,2]+size))
    )
    numpoints[i]=len(points[cube])

（循环是对单个立方体排序，“立方体”创建一个索引的布尔数组。）然后我将点[cube]存储在某个地方，但这并不是减慢速度的原因，而是“cube=”的创建。在

我想加速这个循环（在MacBookPro上完成需要几十秒）。我尝试重写C中的“cube=”部分，如下所示：

^{pr2}$

这使它的速度加快了一点，而不是两倍。在C中重写两个循环实际上只比原来的numpy版本快一点，因为经常引用数组对象来跟踪哪些点在哪个立方体中。我怀疑这样做有可能快得多，而且我遗漏了一些东西。有人能建议如何加快速度吗？？我是numpy/python的新用户，请提前感谢。在

Answer 1

你可以用scipy.space.cKDTree来加速这种计算。在

代码如下：

import time
import numpy as np

#### create some sample data ####
np.random.seed(1)

V_NUM = 6000
P_NUM = 8000

size = 0.1

vertices = np.random.rand(V_NUM, 3)
points = np.random.rand(P_NUM, 3)

numpoints = np.zeros(V_NUM, np.int32)

#### brute force ####
start = time.clock()
for i in xrange(len(vertices)):        
    cube=((vertices[i,0]<= points[:,0]) & 
    (points[:,0]<(vertices[i,0]+size)) & 
    (vertices[i,1]<= points[:,1]) & 
    (points[:,1] < (vertices[i,1]+size)) &
    (vertices[i,2]<= points[:,2]) & 
    (points[:,2] < (vertices[i,2]+size))
    )
    numpoints[i]=len(points[cube])

print time.clock() - start

#### KDTree ####
from scipy.spatial import cKDTree
center_vertices = vertices + [[size/2, size/2, size/2]]
start = time.clock()
tree_points = cKDTree(points)
_, result = tree_points.query(center_vertices, k=100, p = np.inf, distance_upper_bound=size/2)
numpoints2 = np.zeros(V_NUM, np.int32)
for i, neighbors in enumerate(result):
    numpoints2[i] = np.sum(neighbors!=P_NUM)

print time.clock() - start
print np.all(numpoints == numpoints2)

• 首先将立方体角点位置更改为中心位置。在

center_vertices = vertices + [[size/2, size/2, size/2]]

• 从点创建cKDTree

tree_points = cKDTree(points)

• 做查询，k是要返回的最近邻数，p=np.inf公司表示最大坐标差距离，距离上界为最大距离。在

_, result = tree_points.query(center_vertices, k=100, p = np.inf, distance_upper_bound=size/2)

输出为：

^{pr2}$

如果一个立方体中有超过100个点，可以通过for循环中的neighbors[-1] == P_NUM来检查这一点，并对这些顶点执行k=1000查询。在