Question

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

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

但是在这一点上，它比Maya成为更多的python问题，因为我的主要瓶颈是花在寻找顶点匹配上的时间。

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

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

Here我发现了有关KDTree和不同算法的信息，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实现进行了快速测试之后，当在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，必须将树pickle d发送到子流程，因为将参数发送到子流程时会自动完成（这似乎仅对cKDTree有效，而对{{1}无效） }），否则必须在每个过程中从头开始创建树。

以下代码显示了具有多进程KDTree与单个进程的酸洗变体。

cKDTree

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

2 个答案: