Multiply 2dim numpy array with vector to get 3dim array

Question

Suppose I have a numpy array which is 2d and an 1d array

In [127]: A = np.array([[1, 2],[3, 4]])

In [128]: B = np.array([10, 100])

What I would like to achieve is to get a 3d array C , where C[:, :, 0] = A*B[0] and C[:, :, 1] = A*B[1] . I was able to do it via np.einsum but looks like an overkill.

In [129]: np.einsum('ij, k -> ijk', A, B)[:, :, 0]
Out[129]: 
array([[10, 20],
       [30, 40]])

In [130]: np.einsum('ij, k -> ijk', A, B)[:, :, 1]
Out[130]: 
array([[100, 200],
       [300, 400]])

Is there a simpler version?

Answer 1

To use B for scaling along the first axis, we can simply use broadcasting like so -

B[:,None,None]*A # with einsum : np.einsum('ij, k -> kij', A, B)

To get the equivalent of C[:, :, 0] = A*B[0] and C[:, :, 1] = A*B[1] , we need to extend A instead -

A[...,None]*B # with einsum : np.einsum('ij, k -> ijk', A, B)

With no sum-reductions happening in here, broadcasting based ones would be faster than einsum ones.