Vector hóa cho các vòng NumPy

Question

Tôi tương đối mới đối với Python và tôi đã có một vòng lặp lồng nhau. Vì các vòng lặp mất một lúc để chạy, tôi đang cố gắng tìm ra cách để vector hóa mã này để nó có thể chạy nhanh hơn.

Trong trường hợp này, toạ độ là mảng 3 chiều trong đó tọa độ [x, 0, 0] và toạ độ [x, 0, 1] là số nguyên và tọa độ [x, 0, 2] là 0 hoặc 1. H là một ma trận thưa thớt SciPy và x_dist, y_dist, z_dist và a là tất cả các phao.

# x_dist, y_dist, and z_dist are floats 
# coord is a num x 1 x 3 numpy array where num can go into the hundreds of thousands 
num = coord.shape[0]  
H = sparse.lil_matrix((num, num)) 
for i in xrange(num): 
    for j in xrange(num): 
     if (np.absolute(coord[i, 0, 0] - coord[j, 0, 0]) <= 2 and 
       (np.absolute(coord[i, 0, 1] - coord[j, 0, 1]) <= 1)): 

      x = ((coord[i, 0, 0] * x_dist + coord[i, 0, 2] * z_dist) - 
       (coord[j, 0, 0] * x_dist + coord[j, 0, 2] * z_dist)) 

      y = (coord[i, 0, 1] * y_dist) - (coord[j, 0, 1] * y_dist) 

      if a - 0.5 <= np.sqrt(x ** 2 + y ** 2) <= a + 0.5: 
       H[i, j] = -2.7 

Tôi cũng đọc phát thanh đó bằng NumPy, nhanh hơn nhiều, có thể dẫn đến lượng bộ nhớ lớn sử dụng từ mảng tạm thời. Nó sẽ là tốt hơn để đi các tuyến đường vectorization hoặc thử và sử dụng một cái gì đó như Cython?

Answer 1

Tính chất của phép tính làm cho việc vectơ hóa với các phương pháp gọn gàng khiến tôi khó hiểu. Tôi nghĩ rằng giải pháp tốt nhất về tốc độ và sử dụng bộ nhớ sẽ là cython. Tuy nhiên, bạn có thể tăng tốc bằng cách sử dụng numba. Dưới đây là một ví dụ (lưu ý rằng thông thường bạn sử dụng autojit như một trang trí):

import numpy as np 
from scipy import sparse 
import time 
from numba.decorators import autojit 
x_dist=.5 
y_dist = .5 
z_dist = .4 
a = .6 
coord = np.random.normal(size=(1000,1000,1000)) 

def run(coord, x_dist,y_dist, z_dist, a): 
    num = coord.shape[0]  
    H = sparse.lil_matrix((num, num)) 
    for i in xrange(num): 
     for j in xrange(num): 
      if (np.absolute(coord[i, 0, 0] - coord[j, 0, 0]) <= 2 and 
        (np.absolute(coord[i, 0, 1] - coord[j, 0, 1]) <= 1)): 

       x = ((coord[i, 0, 0] * x_dist + coord[i, 0, 2] * z_dist) - 
        (coord[j, 0, 0] * x_dist + coord[j, 0, 2] * z_dist)) 

       y = (coord[i, 0, 1] * y_dist) - (coord[j, 0, 1] * y_dist) 

       if a - 0.5 <= np.sqrt(x ** 2 + y ** 2) <= a + 0.5: 
        H[i, j] = -2.7 
    return H 

runaj = autojit(run) 

t0 = time.time() 
run(coord,x_dist,y_dist, z_dist, a) 
t1 = time.time() 
print 'First Original Runtime:', t1 - t0 

t0 = time.time() 
run(coord,x_dist,y_dist, z_dist, a) 
t1 = time.time() 
print 'Second Original Runtime:', t1 - t0 

t0 = time.time() 
run(coord,x_dist,y_dist, z_dist, a) 
t1 = time.time() 
print 'Third Original Runtime:', t1 - t0 

t0 = time.time() 
runaj(coord,x_dist,y_dist, z_dist, a) 
t1 = time.time() 
print 'First Numba Runtime:', t1 - t0 

t0 = time.time() 
runaj(coord,x_dist,y_dist, z_dist, a) 
t1 = time.time() 
print 'Second Numba Runtime:', t1 - t0 

t0 = time.time() 
runaj(coord,x_dist,y_dist, z_dist, a) 
t1 = time.time() 
print 'Third Numba Runtime:', t1 - t0 

tôi nhận được kết quả này:

First Original Runtime: 21.3574919701 
Second Original Runtime: 15.7615520954 
Third Original Runtime: 15.3634860516 
First Numba Runtime: 9.87108802795 
Second Numba Runtime: 9.32944011688 
Third Numba Runtime: 9.32300305367 

Answer 2

Đây là cách tôi sẽ vectorize mã của bạn, một số cuộc thảo luận về những hãy cẩn thận sau :

import numpy as np 
import scipy.sparse as sps 

idx = ((np.abs(coord[:, 0, 0] - coord[:, 0, 0, None]) <= 2) & 
     (np.abs(coord[:, 0, 1] - coord[:, 0, 1, None]) <= 1)) 

rows, cols = np.nonzero(idx) 
x = ((coord[rows, 0, 0]-coord[cols, 0, 0]) * x_dist + 
    (coord[rows, 0, 2]-coord[cols, 0, 2]) * z_dist) 
y = (coord[rows, 0, 1]-coord[cols, 0, 1]) * y_dist 
r2 = x*x + y*y 

idx = ((a - 0.5)**2 <= r2) & (r2 <= (a + 0.5)**2) 

rows, cols = rows[idx], cols[idx] 
data = np.repeat(2.7, len(rows)) 

H = sps.coo_matrix((data, (rows, cols)), shape=(num, num)).tolil() 

như bạn ghi chú khác, các vấn đề sẽ đi kèm với idx mảng đầu tiên, vì nó sẽ có hình dạng (num, num), vì vậy nó wil l có thể thổi bộ nhớ của bạn thành từng mảnh nếu num là "thành hàng trăm nghìn".

Một giải pháp tiềm năng là chia nhỏ vấn đề của bạn thành các phần có thể quản lý được. Nếu bạn có một mảng phần tử 100.000, bạn có thể chia nó thành 100 phần của 1.000 phần tử và chạy phiên bản đã sửa đổi của mã ở trên cho mỗi kết hợp của 10.000 khối. Bạn sẽ chỉ cần một phần tử 1.000.000 idx mảng (mà bạn có thể phân bổ trước và sử dụng lại để có hiệu suất tốt hơn), và bạn sẽ có một vòng lặp chỉ 10.000 lần lặp lại, thay vì 10.000.000.000 thực hiện hiện tại của bạn. Nó là một chương trình song song của người nghèo, mà bạn thực sự có thể cải thiện bằng cách có một số các khối được xử lý song song nếu bạn có một máy đa lõi.