I have two matrices A,B
, and I'd like to create a 3-d array C
such that
C[k,i,j]=A[k,j]*B[k,i]
I was thinking about using np.einsum
but couldn't find a way and I'm not sure it's useful here.
A naïve loop is possible, but sounds quite inefficient.
UPDATE
C=np.einsum('kj,ki->kij',A,B)
works and is quite elegant, this is on top of @Warren Weckesser 's answer
CodePudding user response:
Basic broadcasting will work, e.g.
In [55]: A # shape is (2, 3)
Out[55]:
array([[4, 9, 5],
[2, 0, 9]])
In [56]: B # shape is (2, 4)
Out[56]:
array([[8, 8, 0, 3],
[6, 8, 3, 8]])
A[:,None,:]
has shape (2, 1, 3), and B[:,:,None]
has shape (2, 4, 1). With broadcasting, the product of those expressions will have shape (2, 4, 3).
In [57]: C = A[:,None,:] * B[:,:,None]
In [58]: C
Out[58]:
array([[[32, 72, 40],
[32, 72, 40],
[ 0, 0, 0],
[12, 27, 15]],
[[12, 0, 54],
[16, 0, 72],
[ 6, 0, 27],
[16, 0, 72]]])
In [59]: C.shape
Out[59]: (2, 4, 3)