CODEWITHSUNDEEP
pythonnumpy
carl
carl

Reputation: 4426

multiply 3D array with 1D array in numpy

I have a 3D array and I have three 1D array, which represents the x-axis, y-axis, and z-axis values. I would like to multiply the 3D array with the 1D arrays. The correct value can be obtained with

   array_x = np.array([1,2,3])
   array_y = np.array([1,2,3])
   array_z = np.array([1,2,3])
   array3D = something
   for ix, x in enumerate(array_x):
       for iy, y in enumerate(array_y):
           for iz, z in enumerate(array_z):
               array3D[ix][iy][iz] *= x*y*z

What is the fastest way to do this in python? I would also like to avoid turning the three 1D arrays into 3D arrays since I need to keep the memory usage low.

Upvotes: 0

Views: 494

Answers (1)

Sayandip Dutta
Sayandip Dutta

Reputation: 15872

You can extend each array along the dimension it is to be multiplied in:

>>> array3D = np.arange(27).reshape((3,3,3))
>>> (array3D*array_x)*array_y[:,None])*array_z[:, None, None]
# or using Einstein Summation convention (np.einsum)
# np.einsum('i, j, k, ijk -> ijk',array_x, array_y, array_z, array3D)

array([[[  0,   2,   6],
        [  6,  16,  30],
        [ 18,  42,  72]],

       [[ 18,  40,  66],
        [ 48, 104, 168],
        [ 90, 192, 306]],

       [[ 54, 114, 180],
        [126, 264, 414],
        [216, 450, 702]]])

Which is the same answer as your nested loop code, but much faster.

Upvotes: 3

Related Questions