是否存在numpy/scipy点积,仅计算结果的对角线条目？

Question

想象一下有2个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)

是否有这样的预制功能,这是否可以有效地完成而无需分配中间(nxn)阵列？

Answer 1

我想我自己得到了它,但仍然会分享解决方案:

因为只获得矩阵乘法的对角线

> Z = N.diag(X.dot(Y))

等价于X行和Y列的标量积的个别和,前一个语句相当于:

> Z = (X * Y.T).sum(-1)

对于原始变量,这意味着:

> result = (A.dot(B) * A).sum(-1)

如果我错了请纠正我,但这应该是......

Answer 2

你几乎可以得到任何你梦寐以求的东西numpy.einsum.直到你开始掌握它,它基本上看起来像黑巫毒...

>>> 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