Использование деревьев для ускорения поиска ближайшего соседа в 3D-массиве с периодическими граничными условиями

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

import numpy as np
from multiprocessing.dummy import Pool as ThreadPool

N = 10000  # Num of objects
# create random point positions
coords = np.random.random((3, N)).transpose()

def NN(point):
    dist = np.abs(np.subtract(coords, point))  # Distance between point and all neighbors xyz values
    dist = np.where(dist > 0.5, dist - 1, dist)  # checking if distance is closer if it wraps around
    return (np.square(dist)).sum(axis=-1).argsort()[1]  # Calc distance and find index of nearest neighbour

# multi threading for speed increase 
pool = ThreadPool(12)
match = pool.map(NN, coords)
pool.close()
pool.join()

Для N ~ 50000 он, как и ожидалось, становится очень медленным.

Я хотел бы знать, как это реализовать, используя такие деревья, как sklearn.BallTree или scipy.spacial.cKDTree и хотел бы сделать это, не дублируя пространство еще 8 раз, как предлагается здесь.


python numpy scipy scikit-learn nearest-neighbor
person JReed    schedule 02.11.2017    source источник

Ответы (1)


arrow_upward
1
arrow_downward

Ни sklearn.BallTree, ни scipy.spatial.cKDTree нельзя легко изменить для обработки периодических граничных условий. Периодические границы требуют принципиально иного набора допущений, чем те, которые используются в этих реализациях.

Вам следует рассмотреть возможность использования альтернативной реализации, например period_kdtree, хотя обратите внимание, что это (несколько устаревший) Реализация Python и будет не такой быстрой, как реализации Cython / C ++, которые вы упомянули.

person jakevdp    schedule 02.11.2017