効率的に1xMを計算するMxM Times Mx1、N Times

Question

N個の行に対してドット積を計算する必要があります。各行に対して、1xM倍MxM倍Mx1を計算する必要があるとします。私が1つの行だけを見ていたら、この計算から1x1を得ました。しかし、私はN行を持っているので、私は行を積み重ね、ドットプロダクトに供給すると思った。しかし、結果としてNxN行列が得られます。私が必要とする結果は対角線上にありますが、スペースを無駄にしないこの計算を行うより速い方法がありますか？理想的には、結果としてNxN行列ではなく、Nx1ベクトルで終わりたいと思います。

例

シングル行

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