如何有效地將numpy數組相加後再乘以數組？

Question

其實我需要計算：

S_i = sum(U_j * U_j.transpose) * K_i 

其中

U_j is a n * k dim matrix, 
K_i is a n * n dim matrix, 
j != i, 
i = 1, 2, ..., n 

而且我用這樣的循環：

import numpy as np 
for i in xrange(n): 
    temp = np.zeros((n, n)) 
    for j in xrange (n): 
     if j != i: 
      temp += np.dot(U[j], U[j].T) 
    S[i] = np.dot(temp, K[i]) 

有沒有更有效的方法來做到這一點？

Answer 1

import numpy as np 

n, k = 30, 40 

U = np.random.random((n, n, k)) 
K = np.random.random((n, n, n)) 

def using_loops(U, K): 
    S = np.empty((n, n, n)) 
    for i in xrange(n): 
     temp = np.zeros((n, n)) 
     for j in xrange (n): 
      if j != i: 
       temp += np.dot(U[j], U[j].T) 
     S[i] = np.dot(temp, K[i]) 
    return S 

def using_einsum(U, K): 
    uut = np.einsum('ijk,ilk->ijl', U, U) 
    total = uut.sum(axis=0) 
    total = total - uut 
    S = np.einsum('ijk,ikl->ijl', total, K) 
    return S 

此測試一些using_loops和using_einsum產生相同的結果。

In [260]: np.allclose(using_loops(U, K), using_einsum(U, K)) 
Out[260]: True 

這表明using_einsum更快;快多少取決於n和k大小：

In [262]: %timeit using_loops(U, K) 
100 loops, best of 3: 17.1 ms per loop 

In [263]: %timeit using_einsum(U, K) 
1000 loops, best of 3: 1.92 ms per loop 

---

在一般情況下，當你看到產品的總和，有一個很好的機會， np.einsum 會產生一個比較快的方法結果。它幾乎肯定會打敗 Python for-loops。