3

In NumPy:

A = np.array([[1,2,3],[4,5,6]])
array([[1, 3, 5],
       [2, 4, 6]])

B = np.array([[1,2],[3,4],[5,6]])
array([[1, 2],
       [3, 4],
       [5, 6]])

A.dot(B)
array([[35, 44],
       [44, 56]])

I only care about getting A.dot(B).diagonal() = array([35, 56])

Is there a way I can get array([35, 56]) without having to compute the inner products of all the rows and columns? I.e. the inner product of the ith row with ith column?

I ask because the performance difference becomes more significant for larger matrices.

BSMP
  • 4,596
  • 8
  • 33
  • 44
ejang
  • 3,982
  • 8
  • 44
  • 70

2 Answers2

6

This is just matrix multiplication for 2D arrays:

C[i, j] = sum(A[i, ] * B[, j])

So since you just want the diagonal elements, looks like you're after

sum(A[i, ] * B[, i]) # for each i

So you could just use list comprehension:

[np.dot(A[i,:], B[:, i]) for i in xrange(A.shape[0])]
# [22, 64]

OR, (and this only works because you want a diagonal so this assumes that if A's dimensions are n x m, B's dimensions will be m x n):

np.sum(A * B.T, axis=1)
# array([22, 64])

(no fancy numpy tricks going on here, just playing around with the maths).

mathematical.coffee
  • 55,977
  • 11
  • 154
  • 194
  • thanks again... so it's not as elegant as http://stackoverflow.com/questions/2301046/compute-only-diagonals-in-matrix-product-in-octave :( – ejang Nov 20 '12 at 01:05
  • Well if you found that question you can translate it directly into numpy (I thought of the exact same thing as soon as I posted the answer and added it to my answer) – mathematical.coffee Nov 20 '12 at 01:07
  • Hi! This worked like a charm, but would you have the math demonstration for `np.diag(A@B) = np.sum(A*B.T , axis=1)` ? – jeannej Sep 11 '19 at 22:38
  • In case some are interested in the math demonstration, I wrote a question/answer [here on math StackExchange](https://math.stackexchange.com/a/3355371/660580) – jeannej Sep 13 '19 at 16:08
0

Can you simply leave out the row in the parameter you don't care about?

The 2x3 x 3x2 gives you a 2x2 result.

A 1x3 x 3x2 matrices will give you only the top row of [A][B], a 1x2 matrix.

EDIT: misread the question. Still, each value in the matrix is produced by the product of the transpose of a column and a row.

nair.ashvin
  • 791
  • 3
  • 11