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

Posted 2023-03-12

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

【英文标题】：Using trees to speed up nearest neighbour search on 3D array with periodic boundary conditions

【发布时间】：2018-04-15 07:25:41

【问题描述】：

使用根据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 次的情况下做到这一点。

【参考方案1】：

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

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