有效地计算1xM乘以MxM乘以Mx1 N次

Question

我需要计算N行的点积;比如说对于每一行我需要计算1xM乘以MxM乘以Mx1。如果我只看一行，从这个计算得到1x1。然而我有N行，所以我想把这些行堆叠起来，然后把它输入点积;但是我得到了一个NxN矩阵。我需要的结果在对角线上，但是有没有一种更快的方法来做这个计算，而且不浪费空间?理想情况下，我希望最终得到一个Nx1向量，而不是NxN矩阵。

例子

单行

r = np.array([[1,2]]).T
R = np.array([[2,2],[2,2]])

为

[[18]]

多个行

rs = np.array([[1,2],[4,4]]).T
R = np.array([[2,2],[2,2]])
print np.dot(np.dot(rs.T,R), rs)
[[ 18  48]
 [ 48 128]]

Answer 1

将np.einsum与内部点积一起使用，如-

np.einsum('ij,ji->i',np.dot(rs.T,R),rs)

示例运行-

In [215]: rs = np.random.rand(3,4)
     ...: R = np.random.rand(3,3)
     ...: out = np.dot(np.dot(rs.T,R), rs)
     ...: 
In [216]: np.diag(out)  # Diagonal elems is the expected o/p
Out[216]: array([ 1.11476081,  1.05112902,  0.32136029,  0.31318894])
In [217]: np.einsum('ij,ji->i',np.dot(rs.T,R),rs)
Out[217]: array([ 1.11476081,  1.05112902,  0.32136029,  0.31318894])

运行测试-

In [233]: rs = np.random.rand(300,400)
In [234]: R = np.random.rand(300,300)
In [235]: %timeit np.diag(np.dot(np.dot(rs.T,R), rs))# Original soln
10 loops, best of 3: 84 ms per loop
In [236]: %timeit np.einsum('ij,kj,ki->j', rs, rs, R)# @DSM's soln
10 loops, best of 3: 65.2 ms per loop
In [237]: %timeit np.einsum('ij,ji->i',np.dot(rs.T,R),rs)
10 loops, best of 3: 37.2 ms per loop