Как эффективно суммировать массивы 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.