Dan Jackson
Dan Jackson

Reputation: 247

Multiplying nth column in 2D array by nth array in 3D array using numpy

I have a 2D and a 3D numpy array and would like to multiply each column of the 2D array by its respective array. Eg multiplying

[[[1. 1.]
  [1. 1.]
  [1. 1.]]

 [[1. 1.]
  [1. 1.]
  [1. 1.]]]

by

[[ 5  6]
 [ 4  7]
 [ 8 10]]

gives

[[[ 5.  5.]
  [ 4.  4.]
  [ 8.  8.]]

 [[ 6.  6.]
  [ 7.  7.]
  [10. 10.]]]

My code currently is:

three_d_array = np.ones([2,3,2])
two_d_array = np.array([(5,6), (4,7), (8,10)])

list_of_arrays = []

for i in range(np.shape(two_d_array)[1]):
    mult = np.einsum('ij, i -> ij', three_d_array[i], two_d_array[:,i])
    list_of_arrays.append(mult)

stacked_array = np.stack(list_of_arrays, 0)

using an answer from Multiplying across in a numpy array but is there a way of doing it without a for loop? Thanks a lot, Dan

Upvotes: 2

Views: 179

Answers (1)

Divakar
Divakar

Reputation: 221514

That nth column in 2D array would be the second axis and by nth array in 3D array, it seems you meant the 2D slice along the first axis. So, the idea would be to align the first axis along three_d_array and second axis along two_d_array. Out of the remaining axes, the first axis from two_d_array seems to be aligned to the second one off three_d_array.

So, to solve it we could use two methods and functions.

Approach #1

Transpose 2D array and then extend dimensions to 3D to have one singleton one at the end and then perform elementwise multiplication with other 3D array, leveraging broadcasting for a vectorized solution -

three_d_array*two_d_array.T[...,None]

Approach #2

With np.einsum -

np.einsum('ijk,ji->ijk',three_d_array, two_d_array)

Upvotes: 2

Related Questions