pythonnumpymultidimensional-arrayelementwise-operations

Element-wise multidimension matrix multiplication with broadcasting in python


Is there a efficient (numpy function) way to do element-wise matrix multiplication of two different-sized arrays that will broadcast into a new array. e.g.,

a = np.arange(24).reshape((2,12)) #gives a 2x12 array 
b = np.arange(36).reshape((3,12)) #gives a 3x12 array

It would then multiply along the dimension with 12 (no summing) to give a final 3-dimensional matrix "c" with shape of 2x3x12 where

c[0,0,:] = [a[0,0]*b[0,0], a[0,1]*b[0,1], ... a[0,11]*b[0,11]]
c[1,0,:] = [a[1,0]*b[0,0], a[1,1]*b[0,1], ... a[1,11]*b[0,11]]
c[0,1,:] = [a[0,0]*b[1,0], a[0,1]*b[1,1], ... a[0,11]*b[1,11]]

I can get what I want through:

a2 = np.repeat(a[..., np.newaxis], 3, axis=2)
b2 = b.T[np.newaxis, ...]
c = np.swapaxes(a2*b2, 1, 2)

that outputs:

array([[[  0,   1,   4,   9,  16,  25,  36,  49,  64,  81, 100, 121],
        [  0,  13,  28,  45,  64,  85, 108, 133, 160, 189, 220, 253],
        [  0,  25,  52,  81, 112, 145, 180, 217, 256, 297, 340, 385]],

       [[  0,  13,  28,  45,  64,  85, 108, 133, 160, 189, 220, 253],
        [144, 169, 196, 225, 256, 289, 324, 361, 400, 441, 484, 529],
        [288, 325, 364, 405, 448, 493, 540, 589, 640, 693, 748, 805]]])

but this feels really inefficient using five-ish numpy commands for something I think shouldn't be too uncommon.


Solution

  • Certainly, you can use broadcasting:

    out = a[:, None, :] * b[None, :, :]
    

    out:

    array([[[  0,   1,   4,   9,  16,  25,  36,  49,  64,  81, 100, 121],
            [  0,  13,  28,  45,  64,  85, 108, 133, 160, 189, 220, 253],
            [  0,  25,  52,  81, 112, 145, 180, 217, 256, 297, 340, 385]],
    
           [[  0,  13,  28,  45,  64,  85, 108, 133, 160, 189, 220, 253],
            [144, 169, 196, 225, 256, 289, 324, 361, 400, 441, 484, 529],
            [288, 325, 364, 405, 448, 493, 540, 589, 640, 693, 748, 805]]])