CODEWITHSUNDEEP
pythonnumpymatrix-multiplication
Shweta
Shweta

Reputation: 1161

Numpy matrix multiplication of 2d matrix to give 3d matrix

I have two numpy arrays, like 

A: = array([[0, 1],  
            [2, 3],  
            [4, 5]])   

B = array([[ 6,  7],  
           [ 8,  9],  
           [10, 11]])

For each row of A and B, say Ra and Rb respectively, I want to calculate transpose(Ra)*Rb. So for given value of A and B, i want following answer:

array([[[ 0,  0],  
        [ 6,  7]],  

       [[ 16,  18],  
        [ 24,  27]],  

       [[ 40,  44],  
        [ 50,  55]]])

I have written the following code to do so:

x = np.outer(np.transpose(A[0]), B[0])
for i in range(1,len(A)):
    x = np.append(x,np.outer(np.transpose(A[i]), B[i]),axis=0)

Is there any better way to do this task.

Upvotes: 3

Views: 1467

Answers (1)

Divakar
Divakar

Reputation: 221554

You can use extend dimensions of A and B with np.newaxis/None to bring in broadcasting for a vectorized solution like so -

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

Explanation : np.outer(np.transpose(A[i]), B[i]) basically does elementwise multiplications between a columnar version of A[i] and B[i]. You are repeating this for all rows in A against corresoinding rows in B. Please note that the np.transpose() doesn't seem to make any impact as np.outer takes care of the intended elementwise multiplications.

I would describe these steps in a vectorized language and thus implement, like so -

  1. Extend dimensions of A and B to form 3D shapes for both of them such that we keep axis=0 aligned and keep as axis=0 in both of those extended versions too. Thus, we are left with deciding the last two axes.
  2. To bring in the elementwise multiplications, push axis=1 of A in its original 2D version to axis=1 in its 3D version, thus creating a singleton dimension at axis=2 for extended version of A.
  3. This last singleton dimension of 3D version of A has to align with the elements from axis=1 in original 2D version of B to let broadcasting happen. Thus, extended version of B would have the elements from axis=1 in its 2D version being pushed to axis=2 in its 3D version, thereby creating a singleton dimension for axis=1.

Finally, the extended versions would be : A[...,None] & B[:,None,:], multiplying whom would give us the desired output.

Upvotes: 5

Related Questions