如何用python加速最近的邻居搜索？

Question

我有一个代码，它计算最近的体素(未赋值)到体素(被赋值)。也就是说，我有一个体素数组，很少的体素已经分配了标量(1,2,3,4....etc)值，很少的体素是空的(假设值为'0')。下面的代码查找与未分配的体素最近的赋值体素，并将该体素分配给同一个标量。因此，标量'0‘的体素将被赋予一个值(1或2或3，.)基于最近的体素。下面的代码可以工作，但是需要花费太多的时间。有什么替代办法吗？或者，如果你对如何进一步改进它有任何反馈？

“”#self.voxels是一个三维数字数组“”

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)，那么计算最近的搜索所需的时间就不再有效了。我怎样才能克服这一切？并行计算能否减少时间(我不知道是否可以并行化代码，如果可以，请告诉我)？

潜在的解决办法：

通过在n_jobs查询中添加cKDTree = -1参数，可以大大提高计算时间。

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

我能够在一个13核心CPU上计算一个数组(400,100,100)在不到一个小时内的距离。我尝试了一个处理器，它需要大约18个小时来完成相同的数组。感谢@gsamaras的回答！

Answer 1

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

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

这个包还提供了Ball Tree算法，您可以测试它与kd-树算法，但是我的预感是kd-树会更好(但这同样取决于您的数据，所以研究！)

您也可能希望使用维数约简，这很容易。我们的想法是减少你的维度，这样你的数据包含的信息就少了，这样就可以更快地解决最近的问题了。当然，这里有一个权衡，准确性！

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

评论：

但是，如果您真的想要高性能，那么就不能使用python，可以切换到c++，例如使用CGAL。

Answer 2

您可以切换到近似近邻(ANN)算法，该算法通常利用复杂的散列或邻近图技术来快速索引数据并执行更快的查询。一个例子是Spotify的烦扰。Annoy的自述模型包括一个图表，它显示了近年来各种ANN算法的精度-性能权衡比较。性能最好的算法(在发布此评论时) 恩索在非度量空间库(NMSLIB)下有一个Python实现。