Is er een numpy/scipy dot-product, dat alleen de diagonale invoer van het resultaat berekent?

Stel je voor dat je 2 numpy arrays hebt:

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

Gewoonlijk is p een kleiner getal (p <= 200), terwijl n willekeurig groot kan zijn.

Ik doe het volgende:

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

Zoals je kunt zien, bewaar ik alleen de n diagonale items, maar er is een tussenliggende (n x n) array berekend waaruit alleen de diagonale items worden bewaard.

Ik wens een functie zoals diag_dot(), die alleen de diagonale invoer van het resultaat berekent en niet het volledige geheugen toewijst.

Een resultaat zou zijn:

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

Is er een vooraf gemaakte functionaliteit zoals deze en kan dit efficiënt worden gedaan zonder dat de tussenliggende (n x n) array moet worden toegewezen?

Antwoord 1, autoriteit 100%

Ik denk dat ik het zelf heb, maar ik zal toch de oplossing delen:

omdat alleen de diagonalen van een matrixvermenigvuldiging worden verkregen

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

is gelijk aan de individuele som van het scalaire product van rijen van X en kolommen van Y, de vorige verklaring is gelijk aan:

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

Voor de oorspronkelijke variabelen betekent dit:

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

Corrigeer me als ik het mis heb, maar dit zou het moeten zijn …

Antwoord 2, autoriteit 66%

Je kunt bijna alles krijgen waar je ooit van hebt gedroomd met numpy.einsum. Totdat je het onder de knie hebt, lijkt het eigenlijk op zwarte voodoo…

>>> 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])

BEWERKENJe kunt het hele ding in één keer zien, het is belachelijk…

>>> 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])

BEWERKJe wilt het echter niet te veel op zichzelf laten bepalen… Ter vergelijking ook het antwoord van de OP op zijn eigen vraag toegevoegd.

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

Antwoord 3, autoriteit 3%

Een voetgangersantwoord, dat de constructie van grote tussenliggende arrays vermijdt, is:

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