0

I have the arrays A and B:

>>> import numpy as np

>>> A = np.ones((3,3,2))

>>> B = np.array([
    [0,0],
    [1,1],
    [2,2],
])

I want to multiply each row of B to each slice of A, such that each row of B gets broadcast over each slice of A, that is:

>>> np.array([A_slice*B_row for A_slice, B_row in zip(A, B)])
[[[0. 0.]
  [0. 0.]
  [0. 0.]]

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

 [[2. 2.]
  [2. 2.]
  [2. 2.]]]

I want the most efficient method for doing this, which I believe may be using np.einsum (however if you think using another method is faster, for example the one I mention below, do tell me).

I have attempted the follwing:

>>> np.einsum('ijk,lk->ijk', A, B)
[[[3. 3.]
  [3. 3.]
  [3. 3.]]

 [[3. 3.]
  [3. 3.]
  [3. 3.]]

 [[3. 3.]
  [3. 3.]
  [3. 3.]]]

As you can see this is clearly not the same output as above.

Another solution that I can think of is:

>>> A*B[:,np.newaxis,:].repeat(3, axis=1)
[[[0. 0.]
  [0. 0.]
  [0. 0.]]

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

 [[2. 2.]
  [2. 2.]
  [2. 2.]]]

Which does give the correct output, but I still really want to know how to do it with np.einsum

EDIT: Warren Weckesser has pointed out in the comments that the solution above can be simplified to A*B[:,np.newaxis,:], which is the cleanest solution I have seen yet without using np.einsum.

Naphat Amundsen
  • 1,519
  • 1
  • 6
  • 17

1 Answers1

0

numpy.einsum soution:

C = np.einsum('ijk,jk->jik', A, B)

Using elipsis: 

C = np.einsum('ij...,j...->ji...', A, B)

Output

[[[0. 0.]
  [0. 0.]
  [0. 0.]]

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

 [[2. 2.]
  [2. 2.]
  [2. 2.]]]
b-fg
  • 3,959
  • 2
  • 28
  • 44
  • Thanks for the attempt, this does indeed work for the case that I illustrated, however I tried to "expand" the problem a little: `A = np.ones((4,3,2))` `B = np.array([[0,0],[1,1],[2,2],[3,3]])` and it got an error saying `ValueError: operands could not be broadcast together with remapped shapes [original->remapped]: (4,3,2)->(3,4,2) (4,2)->(4,newaxis,2)` – Naphat Amundsen May 07 '20 at 11:43