【问题标题】：How can I speed up nearest neighbor search with python?如何使用 python 加快最近邻搜索？
【发布时间】：2017-02-01 07:49:05
【问题描述】：

我有一个代码，它计算最近的体素（未分配）到体素（已分配）。也就是说，我有一个体素数组，少数体素已经分配了标量（1,2,3,4....等）值，并且很少有体素是空的（假设值为“0”）。下面的代码找到最接近未分配体素的已分配体素，并为该体素分配相同的标量。因此，标量为“0”的体素将根据最近的体素分配一个值（1 或 2 或 3，...）。下面的代码有效，但需要太多时间。有替代方案吗？或者您对如何进一步改进有任何反馈？

""" #self.voxels 是一个 3D 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) 数组的距离。我尝试使用 1 个处理器，完成相同的阵列大约需要 18 个小时。感谢@gsamaras 的回答！

【问题讨论】：

假设您可以尝试scikit-learn.org/stable/modules/generated/… 的方法，但我认为问题出在内存容量上。 (500,500,500) 真是个巨大的物体。
你好@gsamaras，谢谢你的回答。我用 sklearn 邻居方法做了 tr，但计算时间似乎没有更大的影响。在我接受您的答案作为最后一个答案之前，我一直在等待其他人是否可以提供任何其他答案。由于您的答案与我的要求更接近，因此我会接受您的答案。再次感谢您！
嗯，我看到了@RavirajpurohitPurushottamr，感谢您的意见，祝您好运！

标签： python performance parallel-processing nearest-neighbor kdtree

【解决方案1】：

您可以切换到近似最近邻 (ANN) 算法，该算法通常利用复杂的散列或邻近图技术来快速索引您的数据并执行更快的查询。一个例子是 Spotify 的Annoy。 Annoy 的自述文件包括一个图表，显示了近年来发布的各种 ANN 算法的精度-性能权衡比较。性能最佳的算法（在发布此评论时）hnsw 在Non-Metric Space Library (NMSLIB) 下有一个 Python 实现。

【讨论】：

【解决方案2】：

试试sklearn.neighbors.NearestNeighbors会很有趣，它提供n_jobs参数：

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

这个包还提供了 Ball Tree 算法，您可以测试它与 kd-tree 算法的对比，但是我的预感是 kd-tree 会更好（但这又取决于您的数据，所以研究一下吧！） .

您可能还想使用降维，这很简单。这个想法是您减少尺寸，因此您的数据包含更少的信息，因此可以更快地解决最近邻问题。当然，这里有一个权衡，准确性！

降维可能/将会降低准确性，但可能值得一试。但是，这通常适用于高维空间，你只是在 3D 中。所以我不知道对于您的具体情况是否可以使用sklearn.decomposition.PCA。

备注：

如果你真的想要高性能，你不会用python得到它，你可以切换到c++，例如使用CGAL。

【讨论】：