CODEWITHSUNDEEP
pythonarraysnumpymultidimensional-arraymatrix-multiplication
omegamanda
omegamanda

Reputation: 398

Array of submatrices to be multiplied by array of scalars

A is a 10 x 10 array of scalars, shape=(10,10)

B is a 10 x 10 array of 3 x 3 matrices, shape=(10,10,3,3)

There are 100 scalars in A and 100 3 x 3 matrices in B. I would like for each 3 x 3 matrix in B to be multiplied by a corresponding scalar in A.

I was able to do this with a simple function like this:

def C(i,j):
   return A[i,j]*B[i,j]

but I wondered if it could be done using even simpler numpy-compatible expressions (without a new function).

Upvotes: 2

Views: 430

Answers (1)

Alex Riley
Alex Riley

Reputation: 176820

You can add new axes to A and then multiply with B to get the desired result:

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

This correctly aligns the axes so each scalar in A multiplies the corresponding 3x3 matrix in B.

A smaller example for demonstration:

>>> A = np.arange(1, 5).reshape(2, 2)
>>> B = np.ones((2, 2, 3, 3)
>>> A[:, :, None, None] * B
array([[[[ 1.,  1.,  1.],
         [ 1.,  1.,  1.],
         [ 1.,  1.,  1.]],

        [[ 2.,  2.,  2.],
         [ 2.,  2.,  2.],
         [ 2.,  2.,  2.]]],


       [[[ 3.,  3.,  3.],
         [ 3.,  3.,  3.],
         [ 3.,  3.,  3.]],

        [[ 4.,  4.,  4.],
         [ 4.,  4.,  4.],
         [ 4.,  4.,  4.]]]]

Upvotes: 3

Related Questions