Wektoryzacja dla pętli NumPy

Question

Jestem stosunkowo nowy w Pythonie i mam zagnieżdżoną pętlę for. Ponieważ pętle for trwają trochę dłużej, próbuję znaleźć sposób na wektoryzację tego kodu, aby działał szybciej.

W tym przypadku, coordin to trójwymiarowa tablica, w której coordin [x, 0, 0] i coord [x, 0, 1] to liczby całkowite, a coordin [x, 0, 2] to 0 lub 1. H jest rzadką macierzą SciPy, a x_dist, y_dist, z_dist i a są zmiennymi.

# 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 

Mam również przeczytać, że nadawanie z NumPy, podczas gdy o wiele szybciej, może prowadzić do dużych ilości zużycia pamięci z tymczasowych tablic. Czy byłoby lepiej pójść drogą wektoryzacji lub spróbować użyć czegoś takiego jak Cython?

Answer 1

Charakter obliczeń sprawia, że ​​trudno jest wektoryzować za pomocą metod numpy, które znam. Myślę, że najlepszym rozwiązaniem pod względem prędkości i wykorzystania pamięci byłby cython. Możesz jednak przyspieszyć używając numba. Oto przykład (uwaga, że ​​zwykle używasz autojit jako dekorator):

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 

uzyskać ten wynik:

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

To jak bym wektoryzacji kodu, dyskusję na temat zastrzeżeń później :

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() 

jak można zauważyć, że problemy będą pochodzić z pierwszego idx tablicy, jak to będzie od kształtu (num, num), więc Wil Prawdopodobnie rozbiję twoją pamięć na kawałki, jeśli num jest "na setki tysięcy".

Jednym z potencjalnych rozwiązań jest rozbicie problemu na łatwe do opanowania części. Jeśli masz tablicę 100 000 elementów, możesz podzielić ją na 100 części po 1000 elementów i uruchomić zmodyfikowaną wersję powyższego kodu dla każdej z 10 000 kombinacji porcji. Potrzebna byłaby tylko tablica o wielkości 1 000 000 elementów idx (którą można wstępnie przydzielić i użyć ponownie w celu uzyskania lepszej wydajności), a zamiast 10 000 000 000 aktualnej implementacji miałbyś pętlę z tylko 10 000 iteracji. Jest to rodzaj schematu równoległego ubogiego człowieka, który można faktycznie poprawić, ponieważ kilka z tych porcji jest przetwarzanych równolegle, jeśli masz maszynę wielordzeniową.