Arrays

Elegant method for element-wise multiplication along shared dimensions of two different-sized arrays?

admin
Welcome to Programming Tutorial official website. Today - we are going to cover how to solve / find the solution of this error Elegant method for element-wise multiplication along shared dimensions of two different-sized arrays? on this date .

Assume that I have two arrays, say array A of shape (X,Y) and array B of shape (X,Y,Z). It is my goal to obtain an array C composed of an element-wise multiplication of A and B along Z. I could do this with a for loop, e.g.:

import numpy as np
C = np.zeros(B.shape)
for z in range(Z):
    C[:,:,z] = A*B[:,:,z]

Ideally, however, I would like to avoid the use of a for loop. Do you know of a more elegant way to achieve this?

Answer

You can element-wise multiply with arrays of shape (x, y) and (z, x, y).

For example:

a = np.arange(12).reshape((4, 3))
b = np.ones((4, 3, 2))

b[:, :, 0] = 10

x = a * b # :( error

However, if we do

c = np.ones((2, 4, 3))
c[0, :, :] = 10
y = a * c

This works and we get

array([[[  0.,  10.,  20.],
        [ 30.,  40.,  50.],
        [ 60.,  70.,  80.],
        [ 90., 100., 110.]],

       [[  0.,   1.,   2.],
        [  3.,   4.,   5.],
        [  6.,   7.,   8.],
        [  9.,  10.,  11.]]])

Now, to move an axis of an array, you can use the np.moveaxis() function.

b_moved = np.moveaxis(b, -1, 0) # move last axis of b to first position

x = a * b_moved # In the shape of b_moved

y = np.moveaxis(x, 0, -1) # Move first axis of x to the last position

And this gives

x: shape 2x4x3
array([[[  0.,  10.,  20.],
        [ 30.,  40.,  50.],
        [ 60.,  70.,  80.],
        [ 90., 100., 110.]],

       [[  0.,   1.,   2.],
        [  3.,   4.,   5.],
        [  6.,   7.,   8.],
        [  9.,  10.,  11.]]])

y: shape 4x3x2
array([[[  0.,   0.],
        [ 10.,   1.],
        [ 20.,   2.]],

       [[ 30.,   3.],
        [ 40.,   4.],
        [ 50.,   5.]],

       [[ 60.,   6.],
        [ 70.,   7.],
        [ 80.,   8.]],

       [[ 90.,   9.],
        [100.,  10.],
        [110.,  11.]]])

In summary, you can write a function that does this for you:

def emult3(arr2d, arr3d):
    # make sure shapes are compatible
    assert (arr2d.ndim, arr3d.ndim) == (2, 3)
    assert arr2d.shape == arr3d.shape[0:2] 

    # do the calculations
    x = arr2d * np.moveaxis(arr3d, -1, 0)
    return np.moveaxis(x, 0, -1)