用Python快速计算XMX^T的对角线

Question

我需要计算XMX ^ T的对角线，但不使用for循环，或者换句话说，替换以下for循环：

X = nump.random.randn(10000, 100)
M = numpy.random.rand(100, 100)
out = numpy.zeros(10000)
for n in range(10000):
  out[n] = np.dot(np.dot(X[n, :], M), X[n, :])

我知道我应该使用numpy.einsum，但是我还没有弄清楚如何使用。非常感谢！

- Yasha Taghia

---

抱歉 M = numpy.random.rand(100, 100) - Yasha Taghia

为什么不使用np.dot(np.dot(x[n, :], M), x[n, :]).diagonal()？ - co2y

Answer 1

当然，有一个np.einsum的方法，如下所示 -

np.einsum('ij,ij->i',X.dot(M),X)

这段代码首先通过X.dot(M)的快速矩阵乘法来实现，然后使用np.einsum保留第一个轴并减少第二个轴的维度。

运行时间测试 -

本节将比较到目前为止发布的所有方法来解决这个问题。

In [132]: # Setup input arrays
     ...: X = np.random.randn(10000, 100)
     ...: M = np.random.rand(100, 100)
     ...: 
     ...: def original_app(X,M):
     ...:     out = np.zeros(10000)
     ...:     for n in range(10000):
     ...:       out[n] = np.dot(np.dot(X[n, :], M), X[n, :])
     ...:     return out
     ...: 

In [133]: np.allclose(original_app(X,M),np.einsum('ij,ij->i',X.dot(M),X))
Out[133]: True

In [134]: %timeit original_app(X,M) # Original solution
10 loops, best of 3: 97.8 ms per loop

In [135]: %timeit np.dot(X, np.dot(M,X.T)).trace()# @Colonel Beauvel's solution
1 loops, best of 3: 2.24 s per loop

In [136]: %timeit np.einsum('ij,jk,ik->i', X, M, X) # @hpaulj's solution
1 loops, best of 3: 442 ms per loop

In [137]: %timeit np.einsum('ij,ij->i',X.dot(M),X) # Proposed in this post
10 loops, best of 3: 28.1 ms per loop

- Divakar

Answer 2

这是一个更简单的例子：

M = array([[ 0,  4,  8],
           [ 1,  5,  9],
           [ 2,  6, 10],
           [ 3,  7, 11]])

X = array([[ 0,  4,  8],
           [ 1,  5,  9],
           [ 2,  6, 10],
           [ 3,  7, 11]])

你所寻找的——对角线元素之和——在数学中更常被称为矩阵的迹。你可以通过以下方式，不使用循环，获得矩阵乘积的迹：

In [102]: np.dot(X, np.dot(M,X.T)).trace()
Out[102]: 692

- Colonel Beauvel

Answer 3

In [210]: X=np.arange(12).reshape(4,3)
In [211]: M=np.ones((3,3))

In [212]: out=np.zeros(4)
In [213]: for n in range(4):
    out[n]= np.dot(np.dot(X[n,:],M), X[n,:])
   .....:     

In [214]: out
Out[214]: array([   9.,  144.,  441.,  900.])

一种使用 einsum 的方法：

In [215]: np.einsum('ij,jk,ik->i', X, M, X)
Out[215]: array([   9.,  144.,  441.,  900.])

比较其他的 einsum:

In [218]: timeit np.einsum('ij,jk,ik->i', X, M, X)
100000 loops, best of 3: 8.98 µs per loop

In [219]: timeit np.einsum('ij,ij->i',X.dot(M),X)
100000 loops, best of 3: 11.9 µs per loop

这样会更快一些，但结果可能与您的较大尺寸不同。

einsum可以避免计算许多不必要的值（与对角线或跟踪方法相比）。

类似地使用einsum - 组合Einsum表达式。

- hpaulj

---

有趣的使用 einsum，但看起来对于列出的数据大小，这可能比原始解决方案更慢。我在我的帖子中添加了运行时测试以针对这些大小。 - Divakar