Question

使用根据this问题的答案改编的代码我可以在考虑周期性边界条件的情况下对3D阵列进行强力NN搜索。然后代码返回最近邻居的索引，并为所有邻居执行此操作。

import numpy as np
from multiprocessing.dummy import Pool as ThreadPool

N = 10000  # Num of objects
# create random point positions
coords = np.random.random((3, N)).transpose()

def NN(point):
    dist = np.abs(np.subtract(coords, point))  # Distance between point and all neighbors xyz values
    dist = np.where(dist > 0.5, dist - 1, dist)  # checking if distance is closer if it wraps around
    return (np.square(dist)).sum(axis=-1).argsort()[1]  # Calc distance and find index of nearest neighbour

# multi threading for speed increase 
pool = ThreadPool(12)
match = pool.map(NN, coords)
pool.close()
pool.join()

对于N~50000，它会变得非常缓慢。

我想知道如何使用sklearn.BallTree或scipy.spacial.cKDTree这样的树来实现这一点，并希望这样做，而不会像建议here那样重复8次空间。< / p>

Answer 1

可以轻松修改sklearn.BallTree和scipy.spatial.cKDTree以处理周期性边界条件。周期性边界需要一组与这些实现中使用的假设完全不同的假设。

您应该考虑使用替代实现，例如periodic_kdtree，但请注意这是一个（有点陈旧的）Python实现，并且不会像您提到的Cython / C ++实现那样快。

使用树来加速具有周期性边界条件的3D阵列上的最近邻搜索

1 个答案: