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

Question

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

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

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

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

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

Here I found out about KDTrees and different algorithms and 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 drop-in 替代品

另请注意，.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])

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

Optimize search of closest four elements in two 3D arrays

python

numpy

scipy

kdtree