gboffi
gboffi

Reputation: 25023

broadcasting — multiply a (NxN) array by a (M) array to have a (NxNxM) array

I have two numpy's ndarrays

a.shape # -> (3, 3)
b.shape # -> (5,)

and I want to compute a c array, defined by

c.shape # -> (3, 3, 5)
c[...,i] = a*b[i]

(the exact shape of c is unimportant because I can always transpose its axes).

It seems to me that posting my coding attempts so far would be unnecessarily embarrassing, it's something like many variations on a[:,None,:]*b[None,:] and that's enough said, isn't it?

I suspect that np.einsum() could be the answer, but the syntax of its subscripts command flies above my head...

Upvotes: 1

Views: 154

Answers (2)

Divakar
Divakar

Reputation: 221504

An alternative solution would be with np.outer as basically we are performing outer product there. Now, np.outer expects the inputs to be 1D arrays and as the pre-processing step it flattens the inputs before performing elementwise multiplication. So, the output from np.outer(a,b) would be 2D and we need to reshape it to desired 3D shape. Thus, the final implementation would look something like this -

np.outer(a,b).reshape(a.shape+(-1,))

A more explicit way to do the reshaping would be like so -

np.outer(a,b).reshape(a.shape+b.shape)

Upvotes: 0

val
val

Reputation: 8689

A nice solution is to use the fantastic einsum function:

>>> a = np.array([[1,2,3],[4,5,6],[7,8,9]])
>>> b = np.array([10, 100, 1000, 10000, 100000])
>>> c = np.einsum('ij,k->ijk', a, b)
>>> c[1, 2, 3]
60000
>>> a[1, 2]
6
>>> b[3]
10000

I like it because it makes the transform ij,k -> ijk very explicit. Also possible, and more concise, is to use simple broadcasting:

>>> c = a[:, :, None] * b[None, None, :]
>>> c.shape
(3L, 3L, 5L)

Where you use : to use a dimension, None to broadcast along one.

Upvotes: 4

Related Questions