Эффективно вычислять 1xM Times MxM Times Mx1, N Times

Question

Мне нужно вычислить точечный продукт для N строк; Скажем, для каждой строки мне нужно вычислить 1xM раз MxM раз Mx1. Если бы я смотрел только на одну строку, у меня было бы 1x1 из этого расчета. У меня есть N строк, поэтому я подумал, что я буду складывать строки и передавать их в dot-продукт; Однако в результате получается матрица 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 Alongwith внутренней дот-продукта, например, так -

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