优化两个 3D 数组中最接近的四个元素的搜索

Question

我有两个用 3D 坐标 (x, y, z) 填充的 numpy 数组。对于第一个数组（“目标”数组）的每个点，我需要找到第二个数组（“源”数组）的 4 个最近点。使用不同的方法找到实际结果没有问题，但我想尽可能加快这个过程。

我需要这个，因为我正在开发一个 Maya 工具，它将存储在网格的每个顶点中的信息传输到第二个网格，并且它们可能具有不同数量的顶点。

但在这一点上，它更像是一个 Python 问题而不是 Maya 问题，因为我的主要瓶颈是寻找顶点匹配所花费的时间。

元素的数量可以从几百到几十万不等，我想确保找到加快搜索速度的最佳方法。 我希望我的工具尽可能快，因为它可能会经常使用，而且每次运行时等待几分钟会很烦人。

我找到了一些有用的答案，让我朝着正确的方向前进：

Here 我发现了 KDTrees 和不同的算法，here 我发现了一些关于多线程的有用注意事项。

这里有一些代码模拟了我将使用的场景以及我尝试过的一些解决方案。

import timeit
import numpy as np
from multiprocessing.pool import ThreadPool
from scipy import spatial

# brut Froce
def bruteForce():
    results = []
    for point in sources:
        dists = ((targets - [point]) ** 2).sum(axis=1)  # compute distances
        ndx = dists.argsort()  # indirect sort
        results.append(zip(ndx[:4], dists[ndx[:4]]))
    return results

# Thread Pool Implementation
def threaded():
    def worker(point):
        dists = ((targets - [point]) ** 2).sum(axis=1)  # compute distances
        ndx = dists.argsort()  # indirect sort
        return zip(ndx[:4], dists[ndx[:4]])


    pool = ThreadPool()
    return pool.map(worker, sources)

# KDTree implementation
def kdTree():
    tree = spatial.KDTree(targets, leafsize=50)
    return [tree.query(point, k=4) for point in sources]

# define the number of points for the two arrays
n_targets = 40000  
n_sources = 40000  

#pick some random points
targets = np.random.rand(n_targets, 3) * 100
sources = np.random.rand(n_sources, 3) * 100



print 'KDTree:   %s' % timeit.Timer(lambda: kdTree()).repeat(1, 1)[0]
print 'bruteforce:   %s' % timeit.Timer(lambda: bruteForce()).repeat(1, 1)[0]
print 'threaded:   %s' % timeit.Timer(lambda: threaded()).repeat(1, 1)[0]

我的结果是：

KDTree:       10.724864464  seconds
bruteforce:   211.427750433 seconds
threaded:     47.3280865123 seconds

最有希望的方法是 KDTree。 起初我认为通过使用一些线程将 KDTree 的工作拆分为单独的任务，我可以进一步加快进程。然而，在使用基本的threading.Thread 实现进行快速测试后，在线程中计算 KDTree 时，它​​的性能似乎更差。 阅读this scipy example我可以看到KDTrees并不是真的适合在并行线程中使用，但我并没有真正理解方式。

然后，我想知道是否有任何其他方法可以优化此代码以更快地执行，可能是通过使用多处理或其他某种技巧来并行解析我的数组。

提前感谢您的帮助！

Answer 1

您可以做一件非常简单但非常有效的事情，那就是从 KDTree 切换到 cKDTree。后者是在纯 Python 中实现的第一个的 Cython 替代品。

还要注意.query 是矢量化的，不需要列表理解。

import scipy.spatial as ss

a = np.random.random((40000,3))
b = np.random.random((40000,3))

tree_py = ss.KDTree(a)
tree_cy = ss.cKDTree(a)

timeit(lambda: tree_cy.query(b, k=4), number=10)*100
# 71.06744810007513
timeit(lambda: tree_py.query(b, k=4), number=1)*1000
# 13309.359921026044

所以这几乎是200x 的免费加速。

Answer 2

对于足够多的源点，多处理可能会提高速度。一个关键点是每个子进程必须持有KDTree 的副本。在 Linux（支持 fork）中，如果在构建树后创建子进程，则会自动完成此操作。

对于 Windows，必须将树发送到子进程 pickled，因为它是在向子进程发送参数时自动完成的（这似乎只适用于 cKDTree，但不适用于 KDTree）或者树必须在每个过程中从头开始创建。

以下代码显示了多进程 cKDTree 与单进程的酸洗变体。

import timeit
import numpy as np
from multiprocessing.pool import Pool
from scipy import spatial


# cKDTree implementation
def ckdTree():
    tree = spatial.cKDTree(targets, leafsize=50)
    return [tree.query(point, k=4) for point in sources]


# Initialization to transfer kdtree
def setKdTree(tree):
    global kdtree

    kdtree = tree

# Worker must not be in another function for multiprocessing
def multiprocKd_worker(point):
    return kdtree.query(point, k=4)


# cKDTree process pool implementation
def multiprocCKd():
    tree = spatial.cKDTree(targets, leafsize=50)

    pool = Pool(initializer=setKdTree, initargs=(tree,))
    return pool.map(multiprocKd_worker, sources)


if __name__ == "__main__":
    # define the number of points for the two arrays
    n_targets = 40000
    n_sources = 40000

    #pick some random points
    targets = np.random.rand(n_targets, 3) * 100
    sources = np.random.rand(n_sources, 3) * 100


    print('cKDTree:   %s' % timeit.Timer(lambda: ckdTree()).repeat(1, 1)[0])
    print('multiprocCKd:   %s' % timeit.Timer(lambda: multiprocCKd()).repeat(1, 1)[0])