将它们相乘后如何有效地求和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 循环。