Home > Back-end >  multiplying "across" in two numpy arrays
multiplying "across" in two numpy arrays

Time:10-13

Given two numpy arrays of shape (25, 2), and (2,), one can easily multiply them across:

import numpy as np

a = np.random.rand(2, 25)
b = np.random.rand(2)

(a.T * b).T  # ok, shape (2, 25)

I have a similar situation where b is of shape (2, 4), and I'd like to get the same results as above for all "4" b. The following works,

a = np.random.rand(25, 2)
b = np.random.rand(2, 4)

c = np.moveaxis([a * bb for bb in b.T], -1, 0)  # shape (2, 4, 25)

but I have a hunch that this is possible without moveaxis.

Any ideas?

CodePudding user response:

An alternative with numpy.einsum:

np.einsum('ij,jk->jki', a, b)

Check results are the same:

(np.einsum('ij,jk->jki', a, b) == c).all()
True

CodePudding user response:

In [185]: a = np.random.rand(2, 25)
     ...: b = np.random.rand(2)

The multiplication is possible with broadcasting:

In [186]: a.shape
Out[186]: (2, 25)
In [187]: a.T.shape
Out[187]: (25, 2)
In [189]: (a.T*b).shape
Out[189]: (25, 2)

(25,2) * (2,) => (25,2) * (1,2) => (25,2). The transpose is a moveaxis, changing the result to (2,25)

In your second case.

In [191]: c = np.moveaxis([a * bb for bb in b.T], -1, 0)
In [192]: c.shape
Out[192]: (2, 4, 25)
In [193]: np.array([a * bb for bb in b.T]).shape
Out[193]: (4, 25, 2)

b.T is (4,2), so bb is (2,); with the (25,2) a, produces (25,2) as above. add in the (4,) iteration.

(25,1,2) * (1,4,2) => (25,4,2), which can be transposed to (2,4,25)

In [195]: (a[:,None]*b.T).shape
Out[195]: (25, 4, 2)
In [196]: np.allclose((a[:,None]*b.T).T,c)
Out[196]: True

(2,4,1) * (2,1,25) => (2,4,25)

In [197]: (b[:,:,None] * a.T[:,None]).shape
Out[197]: (2, 4, 25)
In [198]: np.allclose((b[:,:,None] * a.T[:,None]),c)
Out[198]: True
  • Related