Różnica między każdą parą kolumn dwóch numpy tablic (jak to zrobić bardziej efektywnie)?

Question

Próbuję zoptymalizować kod, a profilując zauważyłem, że ta konkretna pętla zajmuje dużo czasu. Czy możesz mi pomóc szybciej to napisać?

import numpy as np 

rows_a, rows_v, cols = (10, 15, 3) 
a = np.arange(rows_a*cols).reshape(rows_a,cols) 
v = np.arange(rows_v*cols).reshape(rows_v,cols) 

c = 0 
for i in range(len(v)): 
    D = ((a-v[i])**2).sum(axis=-1) 
    c += D.min() 

print(c) 

Czy istnieje funkcja numpy, która może to zrobić wydajnie?

Answer 1

import numpy as np 

rows_a, rows_v, cols = (10, 15, 3) 
a = np.arange(rows_a*cols).reshape(rows_a,cols) 
v = np.arange(rows_v*cols).reshape(rows_v,cols) 

def using_loop(): 
    c = 0 
    for i in range(len(v)): 
     D = ((a-v[i])**2).sum(axis=-1) 
     c += D.min() 
    return c 

def using_broadcasting(): 
    return ((a[:,np.newaxis,:]-v)**2).sum(axis=-1).min(axis=0).sum() 

---

In [106]: %timeit using_loop() 
1000 loops, best of 3: 233 µs per loop 

In [107]: %timeit using_broadcasting() 
10000 loops, best of 3: 29.1 µs per loop 

In [108]: assert using_loop() == using_broadcasting() 

---

Przy użyciu NumPy zwykle pomaga wyeliminować for-loops (jeśli to możliwe) oraz ekspresji obliczeniowej operacji przeprowadzonych na całych tablic - lub co najmniej w macierzy, które są tak duże, jak to możliwe. W ten sposób odciążasz więcej pracy do szybkich algorytmów napisanych w języku C lub Fortran bez pośredniego kodu Pythona.

W oryginalnym kodzie ma kształt (10,) dla każdej iteracji pętli. Ponieważ istnieje 15 iteracji pętli, gdybyśmy mogli wyrazić wszystkie wartości dla D ze wszystkich 15 iteracji jednocześnie jako jedną dużą tablicę, wówczas D miałby kształt (10, 15). W rzeczywistości, można to zrobić:

Od a ma kształt (10,3), a[:, np.newaxis, :] ma kształt (10,1,3).

Korzystanie NumPy broadcasting, ponieważ v ma kształt (15,3),

a[:,np.newaxis,:]-v 

ma kształtować (10,15,3). Kwadrowanie, a następnie zsumowanie na ostatniej osi daje tablicę o kształcie (10, 15). Jest to nowy D:

In [109]: ((a[:,np.newaxis,:]-v)**2).sum(axis=-1).shape 
Out[109]: (10, 15) 

Gdy masz D, reszta obliczeń następuje w sposób naturalny.