Python 用树加速周期边界条件下三维阵列的最近邻搜索

使用根据问题答案改编的代码，我可以在考虑周期性边界条件的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，它会像预期的那样变得非常缓慢

---

我想知道我将如何使用诸如或之类的树来实现这一点，并且不想像建议的那样重复空间8次。

无论是

sklearn.BallTree

还是

scipy.spatial.cKDTree

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

您应该考虑使用另一种实现，例如，注意这是一个（有些陈旧的）Python实现，它将不像您提到的Cython /C++实现快。