是否有一个numpy/scipy点积，只计算结果的对角线项？

Question

假设有两个numpy数组：

> A, A.shape = (n,p)
> B, B.shape = (p,p)

通常p是一个较小的数（p<；=200），而n可以任意大。

我正在做以下工作：

result = np.diag(A.dot(B).dot(A.T))

如您所见，我只保留了n个对角线项，但是有一个中间（nxn）数组，从这个数组中只保留了对角线项。

我想要一个像diag_dot（）这样的函数，它只计算结果的对角线项，而不分配完整的内存。

其结果是：

> result = diag_dot(A.dot(B), A.T)

是否有这样一个预先定义的功能，并且可以在不需要分配中间（n x n）数组的情况下高效地完成这项工作？

Answer 1

你几乎可以用^{}得到任何你梦寐以求的东西。直到你开始掌握它的窍门，它基本上看起来像黑色伏都教。。。

>>> a = np.arange(15).reshape(5, 3)
>>> b = np.arange(9).reshape(3, 3)

>>> np.diag(np.dot(np.dot(a, b), a.T))
array([  60,  672, 1932, 3840, 6396])
>>> np.einsum('ij,ji->i', np.dot(a, b), a.T)
array([  60,  672, 1932, 3840, 6396])
>>> np.einsum('ij,ij->i', np.dot(a, b), a)
array([  60,  672, 1932, 3840, 6396])

编辑实际上，你可以在一次拍摄中获得全部内容，这太荒谬了。。。

>>> np.einsum('ij,jk,ki->i', a, b, a.T)
array([  60,  672, 1932, 3840, 6396])
>>> np.einsum('ij,jk,ik->i', a, b, a)
array([  60,  672, 1932, 3840, 6396])

编辑您不想让它自己计算太多。。。还为自己的问题添加了OP的答案以供比较。

n, p = 10000, 200
a = np.random.rand(n, p)
b = np.random.rand(p, p)

In [2]: %timeit np.einsum('ij,jk,ki->i', a, b, a.T)
1 loops, best of 3: 1.3 s per loop

In [3]: %timeit np.einsum('ij,ij->i', np.dot(a, b), a)
10 loops, best of 3: 105 ms per loop

In [4]: %timeit np.diag(np.dot(np.dot(a, b), a.T))
1 loops, best of 3: 5.73 s per loop

In [5]: %timeit (a.dot(b) * a).sum(-1)
10 loops, best of 3: 115 ms per loop