Оптимизируйте поиск ближайших четырех элементов в двух 3D-массивах - PullRequest
Новая
       36

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

0 голосов
/ 24 мая 2019

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

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

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

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

Ответы [ 2 ]

1 голос
/ 24 мая 2019

Существует одна очень простая, но чрезвычайно эффективная вещь, которую вы можете сделать, это переключиться с KDTree на cKDTree. Последний является заменой первого в Cython, который реализован на чистом Python.

Также обратите внимание, что .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 ускорение бесплатно.

0 голосов
/ 25 мая 2019

Для достаточно большого числа исходных точек многопроцессорная обработка может дать выигрыш в скорости.Важным моментом является то, что каждый подпроцесс должен содержать копию KDTree.В Linux (с поддержкой fork) это выполняется автоматически, если создаются подпроцессы после построения дерева.

Для Windows дерево должно быть отправлено pickle d в подпроцессы, как это делается автоматически при отправке параметровв подпроцесс (который, кажется, работает только для 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])
...