Multiplying a 3D numpy array by a 2D numpy array

Question

I need to multiply a 3D numpy array by a 2D numpy array.

Let's say the 3D array A has shape (3, 100, 500) and the 2D array B has shape (3, 100). I need element wise multiplication for each of those 500 axes in the 3D array by the 2D array and then I need to sum along the first axis of the resultant array yielding an array of size (100, 500).

I can get there with a couple of for loops, but surely there must be a numpy function which will achieve this in 1 line? I have had a look at np.tensordot, np.dot, np.matmul, np.prod and np.sum, but none of these functions will do exactly that.

Please provide an example. Just to be clear, are you asking about element-wise multiplication in the first step? — Mad Physicist
– Mad Physicist, Commented Mar 10, 2020 at 14:30

FBruzzesi · Accepted Answer · 2020-03-10 14:31:17Z

2

We can exploit numpy broadcasting:

import numpy as np

a = np.random.rand(3,100,500)
b = np.random.rand(3,100)

# add new axis to b to use numpy broadcasting
b = b[:,:,np.newaxis]
#b.shape = (3,100,1)

# elementwise multiplication
m = a*b
# m.shape = (3,100,500)

# sum over 1st axis
s = np.sum(m, axis=0)

#s.shape = (100,500)

answered Mar 10, 2020 at 14:31

FBruzzesi

6,6143 gold badges19 silver badges42 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Mad Physicist · Accepted Answer · 2020-03-10 14:32:02Z

0

You can broadcast by adding a new unit axis to the 2D array:

np.sum(A * B[..., None], axis=0)

None in an index introduces a unit dimension at that position, which can be used to align axes for broadcasting. ... is shorthand for : as many times as there are dimensions: in this case it's equivalent to :, : since B is 2D.

An alternative way to write it would be

(A * B.reshape(*B.shape, 1)).sum(axis=0)

answered Mar 10, 2020 at 14:32

Mad Physicist

116k29 gold badges202 silver badges292 bronze badges

Comments

Nik P · Accepted Answer · 2020-03-10 14:36:13Z

0

You can try the following which should work,

np.sum(A.T*B.T,axis=-1).T

This will give you shape (100,500)

answered Mar 10, 2020 at 14:36

Nik P

2241 silver badge5 bronze badges

Comments

Sebastian Dziadzio · Accepted Answer · 2020-03-10 14:42:25Z

0

You can easily express these operations using np.einsum, in this case:

np.einsum("ijk,ij->jk", A, B)

edited Mar 10, 2020 at 14:42

answered Mar 10, 2020 at 14:28

Sebastian Dziadzio

5303 silver badges14 bronze badges

Collectives™ on Stack Overflow

Multiplying a 3D numpy array by a 2D numpy array

4 Answers 4

Comments

Comments

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

4 Answers 4

Comments

Comments

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related