Is there a numpy/scipy dot product, calculating only the diagonal entries of the result?

I think i got it on my own, but nevertheless will share the solution:

since getting only the diagonals of a matrix multiplication

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

is equivalent to the individual sum of the scalar product of rows of X and columns of Y, the previous statement is equivalent to:

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

For the original variables this means:

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

Please correct me if I am wrong but this should be it ...

You can get almost anything you ever dreamed of with numpy.einsum. Until you start getting the hang of it, it basically seems like black 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])

EDIT You can actually get the whole thing in a single shot, it's ridiculous...

>>> 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 You don't want to let it figure too much on its own though... Added the OP's answer to its own question for comparison also.

n, p = 10000, 200a = 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 loopIn [3]: %timeit np.einsum('ij,ij->i', np.dot(a, b), a)10 loops, best of 3: 105 ms per loopIn [4]: %timeit np.diag(np.dot(np.dot(a, b), a.T))1 loops, best of 3: 5.73 s per loopIn [5]: %timeit (a.dot(b) * a).sum(-1)10 loops, best of 3: 115 ms per loop

A pedestrian answer, which avoids the construction of large intermediate arrays is:

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