Есть ли продукт numpy/scipy dot, вычисляя только диагональные записи результата?

Представьте, что у вас 2 массива numpy:

> A, A.shape = (n,p)
> B, B.shape = (p,p)

Обычно p является меньшим числом (p <= 200), тогда как n может быть сколь угодно большим.

Я делаю следующее:

result = np.diag(A.dot(B).dot(A.T))

Как вы можете видеть, я сохраняю только n диагональных записей, но есть промежуточный (n x n) массив, рассчитанный, из которого сохраняются только диагональные записи.

Мне нужна функция типа diag_dot(), которая вычисляет только диагональные записи результата и не выделяет полную память.

Результат:

> result = diag_dot(A.dot(B), A.T)

Есть ли готовые функции, подобные этому, и это можно сделать эффективно без необходимости выделения промежуточного (n x n) массива?

Ответ 1

Я думаю, что получил его сам по себе, но, тем не менее, поделится решением:

так как получая только диагонали матричного умножения

> Z = N.diag(X.dot(Y))

эквивалентно индивидуальной сумме скалярного произведения строк X и столбцов Y, предыдущий оператор эквивалентен:

> Z = (X * Y.T).sum(-1)

Для исходных переменных это означает:

> result = (A.dot(B) * A).sum(-1)

Пожалуйста, поправьте меня, если я ошибаюсь, но это должно быть...

Ответ 2

Вы можете получить почти все, что вы когда-либо мечтали с помощью numpy.einsum. Пока вы не начнете ощущать это, это в основном похоже на черный вуду...

>>> a = np.arange(15).reshape(5, 3)
>>> b = np.arange(9).reshape(3, 3)

>>> np.diag(np.dot(np.dot(a, b), a.T))
array([  60,  672, 1932, 3840, 6396])
>>> np.einsum('ij,ji->i', np.dot(a, b), a.T)
array([  60,  672, 1932, 3840, 6396])
>>> np.einsum('ij,ij->i', np.dot(a, b), a)
array([  60,  672, 1932, 3840, 6396])

РЕДАКТИРОВАТЬ Вы можете получить все это за один выстрел, это смешно...

>>> np.einsum('ij,jk,ki->i', a, b, a.T)
array([  60,  672, 1932, 3840, 6396])
>>> np.einsum('ij,jk,ik->i', a, b, a)
array([  60,  672, 1932, 3840, 6396])

РЕДАКТИРОВАТЬ. Вы не хотите, чтобы это слишком сильно задумывалось, хотя... Добавил ответ OP на свой собственный вопрос и для сравнения.

n, p = 10000, 200
a = np.random.rand(n, p)
b = np.random.rand(p, p)

In [2]: %timeit np.einsum('ij,jk,ki->i', a, b, a.T)
1 loops, best of 3: 1.3 s per loop

In [3]: %timeit np.einsum('ij,ij->i', np.dot(a, b), a)
10 loops, best of 3: 105 ms per loop

In [4]: %timeit np.diag(np.dot(np.dot(a, b), a.T))
1 loops, best of 3: 5.73 s per loop

In [5]: %timeit (a.dot(b) * a).sum(-1)
10 loops, best of 3: 115 ms per loop

Ответ 3

Пешеходный ответ, который позволяет избежать создания больших промежуточных массивов:

result=np.empty( [n.], dtype=A.dtype )
for i in xrange(n):
    result[i]=A[i,:].dot(B).dot(A[i,:])