Matrix Multiplication: Multiply each row of matrix by another 2D matrix in Python

Question

I am trying to remove the loop from this matrix multiplication (and learn more about optimizing code in general), and I think I need some form of np.broadcasting or np.einsum, but after reading up on them, I'm still not sure how to use them for my problem.

A = np.array([[1, 2, 3, 4, 5],
         [6, 7, 8, 9, 10],
         [11,12,13,14,15]])
#A is a 3x5 matrix, such that the shape of A is (3, 5) (and A[0] is (5,))

B = np.array([[1,0,0],
         [0,2,0],
         [0,0,3]])
#B is a 3x3 (diagonal) matrix, with a shape of (3, 3)

C = np.zeros(5)
for i in range(5):
    C[i] = np.linalg.multi_dot([A[:,i].T, B, A[:,i]])

#Each row of matrix math is [1x3]*[3x3]*[3x1] to become a scaler value in each row
#C becomes a [5x1] matrix with a shape of (5,)

I know I can't just do np.multidot by itself, because that results in a (5,5) array.

I also found this: Multiply matrix by each row of another matrix in Numpy, but I can't tell if it's actually the same problem as mine.

Answer 1

In [601]: C
Out[601]: array([436., 534., 644., 766., 900.])

This is a natural for einsum. I use i as you do, to denote the index that carries through to the result. j and k are indices that are used in the sum of products.

In [602]: np.einsum('ji,jk,ki->i',A,B,A)
Out[602]: array([436, 534, 644, 766, 900])

It probably can also be done with mutmul, though it may require adding a dimension and latter squeezing.

dot approaches that use diag do a lot more work than necessary. The diag throws out a lot of values.

To use matmul we have to make the i dimension the first of 3d arrays. That's the 'passive' one carries over to the result:

In [603]: A.T[:,None,:]@[email protected][:,:,None]
Out[603]: 
array([[[436]],     # (5,1,1) result

       [[534]],

       [[644]],

       [[766]],

       [[900]]])
In [604]: (A.T[:,None,:]@[email protected][:,:,None]).squeeze()
Out[604]: array([436, 534, 644, 766, 900])

Or index the extra dimensions away: (A.T[:,None,:]@[email protected][:,:,None])[:,0,0]

Answer 2

You can chain to calls to dot together, then get the diagonal:

# your original output:
# >>> C
# array([436., 534., 644., 766., 900.])

>>> np.diag(np.dot(np.dot(A.T,B), A))
array([436, 534, 644, 766, 900])

Or equivalently, use your original multi_dot train of thought, but take the diagonal of the resulting 5x5 array. This may have some performance boosts (according to the docs)

>>> np.diag(np.linalg.multi_dot([A.T, B, A]))
array([436, 534, 644, 766, 900])