Оптимизация поиска ближайших четырех элементов в двух трехмерных массивах

У меня есть два пустых массива, заполненных трехмерными координатами (x, y, z). Для каждой точки первого массива («целевой» массив) мне нужно найти 4 ближайшие точки второго массива («исходный» массив). У меня нет проблем с поиском фактических результатов с использованием различных методов, но я хочу максимально ускорить процесс.

Мне это нужно, потому что я работаю над инструментом Maya, который передает информацию, хранящуюся в каждой вершине сетки, во вторую сетку, и у них может быть разное количество вершин.

На данный момент, однако, это становится больше проблемой Python, чем проблемой Maya, поскольку моим основным узким местом является время, затрачиваемое на поиск совпадений вершин.

Количество элементов может варьироваться от нескольких сотен до сотен тысяч, и я хочу быть уверен, что найду лучший способ ускорить поиск. Я хотел бы, чтобы мой инструмент работал как можно быстрее, так как он может использоваться очень часто, и минуты ожидания каждый раз, когда он должен запускаться, были бы довольно раздражающими.

Я нашел несколько полезных ответов, которые привели меня в правильном направлении:

Здесь я узнал о KDTrees и различных алгоритмах и здесь Я нашел несколько полезных соображений по многопоточности.

Вот некоторый код, имитирующий сценарий, с которым я буду работать, и несколько решений, которые я пробовал.

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 вычисляется в потоке. Читая этот пример scipy, я вижу, что KDTrees не очень подходят для использования в параллельных потоках, но я не очень понял, как.

Тогда мне было интересно, есть ли какой-либо другой способ оптимизировать этот код, чтобы он работал быстрее, возможно, используя многопроцессорность или какой-либо другой трюк для параллельного анализа моих массивов.

Заранее спасибо за помощь!