如何有效地将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。