C'è un prodotto numpy/scipy dot, che calcola solo le voci diagonali del risultato?

Question

immaginare di avere matrici 2 NumPy:

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

Tipicamente p è un numero inferiore (p < = 200), mentre n può essere arbitrariamente grande.

sto facendo quanto segue:

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

Come potete vedere, sto mantenendo solo i n elementi diagonali, tuttavia v'è una (n x n) gamma intermedia calcolata da cui sono conservati solo gli elementi diagonali.

Desidero una funzione come diag_dot(), che calcola solo le voci diagonali del risultato e non alloca la memoria completa.

Un risultato sarebbe:

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

C'è una funzionalità premade come questo e questo può essere fatto in modo efficiente, senza la necessità per l'assegnazione del (n x n) matrice intermedia?

Answer 1

Penso che ho avuto per conto mio, ma comunque si divideranno la soluzione:

poiché ottenere solo le diagonali di una matrice moltiplicazione

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

è equivalente alla somma individuale del prodotto scalare di righe di X e Y colonne, l'istruzione precedente è equivalente a:

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

ai variabili originali questo significa:

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

favore correggetemi se sbaglio, ma questo dovrebbe essere lo ...

Answer 2

È possibile ottenere quasi tutto quello che mai sognato con numpy.einsum. Fino a quando non iniziare a ricevere il blocco di esso, sembra praticamente come voodoo nero ...

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

EDIT Si può effettivamente ottenere il tutto in un solo colpo, è ridicolo ...

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

EDIT Non si vuole lasciare che sia troppo per conto proprio ... Aggiunta la risposta dell'OP alla propria domanda anche per il confronto.

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 

Answer 3

Una risposta pedonale, che evita la costruzione di grandi array intermedi è:

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