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

Question

У меня есть массив точек (баллов), состоящий из значений ~ 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]) 

(Петля заказать отдельные кубики, и «Куб» создает логический массив индексов.) Я тогда магазин точки [куб] где-то , но это не то, что меня замедляет; это создание «куба =».

Я хотел бы ускорить этот цикл (требуется несколько секунд для работы в macbook pro). Я пытался переписывать «куб =» часть в 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/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]]

• Создать cKDTree из точек

tree_points = cKDTree(points)

• ли запрос, к это число ближайших соседей, чтобы вернуться, р = 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 в цикл, и сделать ак = 1000 запросов для этих вершин.