python - 在 python 中快速计算 XMX^T 对角线的方法

Question

我需要在没有 for 循环的情况下计算 XMX^T 的对角线，或者换句话说，替换以下 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，但我一直无法弄清楚如何使用？

非常感谢!

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