如何在将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])

有没有更高效的方法来完成这个任务？

- LittleLittleQ

---

for j in xrange (n) and j != i 应该会抛出一个语法错误。你是不是想要将其拆分成一个 for-loop 和一个 if-statement？ - unutbu

@unutbu 是的，你说得对，我想做的是 for j in xrange(n): if j != i: 但是我在语法上犯了一个错误，我想知道有没有更好的方法来做这件事~ - LittleLittleQ

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

这表明使用_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循环。

- unutbu

---

我已经阅读了有关“np.einsum”使用的文档，但是我对参数“operands”感到非常困惑。最后一个示例是np.einsum('ki,jk->ij', a, b)，我认为如果我没有错的话，它等同于np.dot(a.T, b.T)。但是为什么操作数是“ki，jk->ij”，而不是“ji，kj->ik”？如果您能为我解释一下就好了~谢谢~~ - LittleLittleQ

如果公式 S_i = sum(U_j * U_j.transpose) * K_i 中的 i 不一定等于 n，而是远小于 n，那么使用 np.dot 和 np.einsum 之间的时间效率差异是否会有所改变？ - LittleLittleQ

关于 numpy.einsum 中的操作数，我已经弄清楚了谢谢 :) - LittleLittleQ

很高兴你已经理解了np.einsum中的下标符号——这需要一些时间来适应。对于第二个问题，你是指你只想计算i的少量值（或单个值），远小于n吗？ - unutbu

没错，那正是我的意思。在这种情况下，使用np.dot和np.einsum的效率如何呢？ - LittleLittleQ

如果您只想计算一个 i 值的 S_i，那么我没有找到比使用单个 for 循环更快的方法。 - unutbu