Изготовление срезов точечного вывода без промежуточных продуктов

Question

У меня есть большой набор матриц 3 x 3 (скажем, n) и соответствующие векторы 3 x 1 и хотели бы умножить каждый вектор на соответствующую ему матрицу. Если укладываю матрицы в n х 3 х 3 ndarray называется R и векторы в 3 х nndarray называется v, я могу получить пачку размноженных векторов через,

import numpy as np 
intermediate = np.dot(R, v) 
out = np.diagonal(intermediate, axis1=0, axis2=2) 

Но это очень неэффективно: np.dot производит массив n x 3 x nintermediate, из которого я вручную выбираю 3 x n фрагмент. Могу ли я каким-то образом создать массив 3 x n, не сделав промежуточным n x 3 x n массив?

Answer 1

Расширение на подсказку, представленной @hpaulj: умножение я описывала может быть осуществлено,

out = np.einsum('ijk,ki->ji', R, v) 

Ускорение по подходу в моем вопросе уже 3 порядков для n = 1000 (!):

%timeit d = np.diagonal(np.dot(R, v), axis1=0, axis2=2) 
10 loops, best of 3: 27.8 ms per loop 

%timeit o = np.einsum('ijk,ki->ji', R, v) 
10000 loops, best of 3: 21.9 µs per loop