50

想象一下有 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）数组的情况下有效地完成？

4

3 回答 3

59

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

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

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

等价于 X 的行和 Y 的列的标量积的单个和，前面的语句等价于：

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

对于原始变量，这意味着：

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

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

于 2013-02-07T19:22:38.687 回答
38

你几乎可以得到任何你梦寐以求的东西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
于 2013-02-07T19:26:51.077 回答
2

避免构建大型中间数组的行人答案是：

result=np.empty([n,], dtype=A.dtype )
for i in xrange(n):
    result[i]=A[i,:].dot(B).dot(A[i,:])
于 2013-02-07T20:09:20.823 回答