Ускорьте запрос массива в Numpy / Python

Question

У меня есть массив точек (называемый points ), состоящий из ~ 30000 значений x, y и z. У меня также есть отдельный массив точек (называемых вершины ), около ~ 40000 значений x, y и z. Последний массив индексирует нижние передние левые углы некоторых кубов с длиной стороны размер. Я хотел бы выяснить, какие точки находятся в каких кубах и сколько точек находятся в каждом кубе. Я написал цикл для этого, который работает так:

for i in xrange(len(vertices)):        
    cube=((vertices[i,0]<= points[:,0]) & 
    (points[:,0]<(vertices[i,0]+size)) & 
    (vertices[i,1]<= points[:,1]) & 
    (points[:,1] < (vertices[i,1]+size)) &
    (vertices[i,2]<= points[:,2]) & 
    (points[:,2] < (vertices[i,2]+size))
    )
    numpoints[i]=len(points[cube])

(Цикл - это порядок отдельных кубов, а «куб» создает логический массив индексов.) Затем я где-то храню точки [куб], но это не то, что меня тормозит; это создание "cube =".

Я бы хотел ускорить этот цикл (на MacBook Pro требуются десятки секунд). Я попытался переписать часть "cube =" в C следующим образом:

for i in xrange(len(vertices)):
    cube=zeros(pp, dtype=bool)
    code="""
            for (int j=0; j<pp; ++j){

                if (vertices(i,0)<= points(j,0))
                 if (points(j,0) < (vertices(i,0)+size))
                  if (vertices(i,1)<= points(j,1))
                   if (points(j,1) < (vertices(i,1)+size))
                    if (vertices(i,2)<= points(j,2))
                     if (points(j,2) < (vertices(i,2)+size))
                      cube(j)=1;
            }
        return_val = 1;"""

    weave.inline(code,
    ['vertices', 'points','size','pp','cube', 'i']) 
    numpoints[i]=len(points[cube])

Это ускорило его чуть больше, чем в два раза. Переписывание обоих циклов в C фактически сделало его лишь немного быстрее, чем в оригинальной версии с numpy-only, из-за частых ссылок на объекты массива, необходимые для отслеживания того, какие точки находятся в каких кубах. Я подозреваю, что это можно сделать гораздо быстрее, и я что-то упускаю. Кто-нибудь может подсказать, как это ускорить ?? Я новичок в numpy / python, и заранее спасибо.

Answer 1

Вы можете использовать scipy.spatial.cKDTree для ускорения такого рода расчетов.

Вот код:

import time
import numpy as np

#### create some sample data ####
np.random.seed(1)

V_NUM = 6000
P_NUM = 8000

size = 0.1

vertices = np.random.rand(V_NUM, 3)
points = np.random.rand(P_NUM, 3)

numpoints = np.zeros(V_NUM, np.int32)

#### brute force ####
start = time.clock()
for i in xrange(len(vertices)):        
    cube=((vertices[i,0]<= points[:,0]) & 
    (points[:,0]<(vertices[i,0]+size)) & 
    (vertices[i,1]<= points[:,1]) & 
    (points[:,1] < (vertices[i,1]+size)) &
    (vertices[i,2]<= points[:,2]) & 
    (points[:,2] < (vertices[i,2]+size))
    )
    numpoints[i]=len(points[cube])

print time.clock() - start

#### KDTree ####
from scipy.spatial import cKDTree
center_vertices = vertices + [[size/2, size/2, size/2]]
start = time.clock()
tree_points = cKDTree(points)
_, result = tree_points.query(center_vertices, k=100, p = np.inf, distance_upper_bound=size/2)
numpoints2 = np.zeros(V_NUM, np.int32)
for i, neighbors in enumerate(result):
    numpoints2[i] = np.sum(neighbors!=P_NUM)

print time.clock() - start
print np.all(numpoints == numpoints2)

• Сначала измените угловое положение куба на центральное.

center_vertices = vertices + [[size/2, size/2, size/2]]

• Создание дерева cKD из точек

tree_points = cKDTree(points)

• Выполнить запрос, k - это число ближайших соседей, которые нужно вернуть, p = np.inf означает расстояние с максимальной разницей в координатах, distance_upper_bound - максимальное расстояние.

_, result = tree_points.query(center_vertices, k=100, p = np.inf, distance_upper_bound=size/2)

Вывод:

2.04113164434
0.11087783696
True

Если в кубе более 100 точек, вы можете проверить это с помощью neighbors[-1] == P_NUM в цикле for и выполнить запрос k = 1000 для этих вершин.