如何使用Python加速最近邻搜索？

Question

我有一段代码，可以计算距离（未被分配的）最近的体素到一个已赋值的体素的距离。我有一个体素数组，其中一些体素已经被标量（1、2、3、4等）赋值，而一些体素为空（假设值为“0”）。下面的代码会找到距离未分配的体素最近的已分配的体素，并将该体素分配相同的标量。因此，标量为'0'的体素将根据最近的体素被赋予某个值（1或2或3等）。下面的代码可以运行，但是它需要太长时间。 是否有替代方法？或者您对如何进一步改进它有任何反馈？

"""＃self.voxels是一个三维numpy数组"""

def fill_empty_voxel1(self,argx, argy, argz):
""" where # argx, argy, argz are the voxel location where the voxel is zero"""
    argx1, argy1, argz1 = np.where(self.voxels!=0)   # find the non zero voxels
    a = np.column_stack((argx1, argy1, argz1)) 
    b = np.column_stack((argx, argy, argz))
    tree = cKDTree(a, leafsize=a.shape[0]+1)
    distances, ndx = tree.query(b, k=1, distance_upper_bound= self.mean) # self.mean is a mean radius search value
    argx2, argy2, argz2 = a[ndx][:][:,0],a[ndx][:][:,1],a[ndx][:][:,2]
    self.voxels[argx,argy,argz] = self.voxels[argx2,argy2,argz2] # update the voxel array

示例

这是一个使用小数据集的简单示例:

import numpy as np
from scipy.spatial import cKDTree
import timeit

voxels = np.zeros((10,10,5), dtype=np.uint8)
voxels[1:2,:,:] = 5.
voxels[5:6,:,:] = 2.
voxels[:,3:4,:] = 1.
voxels[:,8:9,:] = 4.
argx, argy, argz = np.where(voxels==0)

tic=timeit.default_timer()
argx1, argy1, argz1 = np.where(voxels!=0)   # non zero voxels
a = np.column_stack((argx1, argy1, argz1)) 
b = np.column_stack((argx, argy, argz))
tree = cKDTree(a, leafsize=a.shape[0]+1)
distances, ndx = tree.query(b, k=1, distance_upper_bound= 5.)
argx2, argy2, argz2 = a[ndx][:][:,0],a[ndx][:][:,1],a[ndx][:][:,2]
voxels[argx,argy,argz] = voxels[argx2,argy2,argz2]
toc=timeit.default_timer()
timetaken = toc - tic #elapsed time in seconds
print '\nTime to fill empty voxels', timetaken

用于可视化：

from mayavi import mlab
data = voxels.astype('float')
scalar_field = mlab.pipeline.scalar_field(data)
iso_surf = mlab.pipeline.iso_surface(scalar_field)
surf = mlab.pipeline.surface(scalar_field)  
vol = mlab.pipeline.volume(scalar_field,vmin=0,vmax=data.max())  
mlab.outline()    
mlab.show()    

现在，如果我的体素数组的尺寸为(500,500,500)，那么计算最近搜索所需的时间就不再高效。我该如何克服这个问题？并行计算能否减少时间（如果您知道是否可以并行化代码，请告诉我）？

一个可能的解决方案是，在cKDTree查询中添加n_jobs = -1参数，可以大幅提高计算速度。

distances, ndx = tree.query(b, k=1, distance_upper_bound= 5., n_jobs=-1)

在一台13核CPU上，我能够在不到一个小时的时间内计算出一个尺寸为(400,100,100)的数组的距离。当我只使用一个处理器时，完成相同的数组需要大约18个小时的时间。 感谢@gsamaras提供的答案！

- Ravi raj purohit Purushottam r

---

作为一种假设，您可以尝试来自http://scikit-learn.org/stable/modules/generated/sklearn.neighbors.NearestNeighbors.html#sklearn.neighbors.NearestNeighbors的方法，但我认为问题在于内存容量。(500,500,500)是一个非常巨大的对象。 - Sergey

你好 @gsamaras，感谢您的回答。我已经使用了sklearn邻居方法进行了tr操作，但是计算时间似乎没有太大影响。在接受您的答案之前，我想等待其他人是否能提供其他答案。由于您的答案更接近我的要求，我将接受您的答案。再次感谢！ - Ravi raj purohit Purushottam r

嗯，我明白了@RavirajpurohitPurushottamr，感谢您的贡献，祝你好运！ - gsamaras

Answer 1

您可以切换到使用近似最近邻（ANN）算法，这些算法通常利用复杂的哈希或接近度图技术快速索引数据并执行更快的查询。其中一个例子是Spotify的Annoy。 Annoy的README中包括一张图表，显示了近年来各种ANN算法的精度-性能折衷比较。在发表此评论时，表现最佳的算法（hnsw）在非度量空间库（NMSLIB）下有一个Python实现。请注意保留HTML标记。

- Mete Kemertas

Answer 2

尝试使用sklearn.neighbors.NearestNeighbors可能会很有趣，它提供了n_jobs参数：

用于运行邻居搜索的并行作业数。

此软件包还提供Ball Tree算法，您可以测试与kd-tree算法相比，但我猜测kd-tree算法会更好（但这也取决于您的数据，因此请进行研究！）。

---

您可能还想使用降维，这很容易。思路是减少维度，因此数据包含的信息更少，从而可以更快地解决最近邻问题。当然，在这里存在一个权衡，即精度！

通过降维，您可能会获得较低的精度，但这值得一试。然而，这通常适用于高维空间，而您只处于三维空间中。因此，我不知道对于您的特定情况是否有意义使用sklearn.decomposition.PCA。

---

一条备注：

如果你真的想要高性能，那么你不能使用Python，你可以转换到C ++，例如使用CGAL。

- gsamaras