How would I achieve this "row in A * all rows in B by col in A" multiplication in NumPy without a loop?

user_14165028 0
Viewed 91

Say I have two matrices, A and B:

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

B = np.array([[0, 1, 0],
              [1, 1, 0],
              [1, 1, 1]])

I want to take one column in A and multiply it by each column in B element-wise, then proceed to the next column in A. So, using just one column as an example, I will use A[:,0] (values 1,2,3), and multiply it by each column in B to get this:

array([[0, 1, 0],
       [2, 2, 0],
       [3, 3, 3]])

I've implemented this using np.einsum like so:

np.einsum('i,ij->ij',A[:,0],B)

I then want to generate a 3D matrix with the depth dimension corresponding to the multiplication by each column in A, which I implemented using a for loop:

np.stack([np.einsum('i,ij->ij',A[:,i],B) for i in range(0,A.shape[1])])

This returns my desired array:

array([[[0, 1, 0],
        [2, 2, 0],
        [3, 3, 3]],

       [[0, 3, 0],
        [2, 2, 0],
        [1, 1, 1]],

       [[0, 2, 0],
        [3, 3, 0],
        [1, 1, 1]]])

How would I go about doing this without the loop? Can this be done purely with np.einsum? Is there another function in NumPy that will do this more simply?

numpy python
3 Answers

Here's a simple way:

A.T[:,:,None]*B

adding the last None in indexing creates a new axis which is then used for broadcasting the elementwise multiplication.

user_45659470

How about this code?

A.T.reshape(3, 3, 1) * B

Reshaping ndarray can make doing many things...

user_131768960

Keeping with your usage of einsum:

np.einsum('ij,ik->jik', A, B)
user_29887300
Related