[FIXED] Prepare numpy array -- similar to broadcasting

Issue

Let's say I have an m-dimensional matrix M and an array N of length n.

The shape of M is (a,b,n,e,f), meaning that on a particular axis (here, 2) the dimensionality of M agrees with that of N.

If I want to multiply them, in this particular case, I would do:

M * N[None, None, :, None, None]

How can I generalize this concept, to cases where the shape of M differs, and the axis may be a different one? I'm sure there must be a convenient function to extend array N accordingly, but I can't find anything

Solution

I'll ilustrate several ways

In [29]: a,b,n,e,f = 2,3,4,5,6
In [30]: M=np.ones((a,b,n,e,f)); N=np.arange(n)

N can be expanded with None; only the trailing ones are required; leading ones are implicit

In [31]: (M*N[:,None,None]).shape
Out[31]: (2, 3, 4, 5, 6)

We could make a indexing tuple with something like:

In [32]: idx = [None]*M.ndim;idx
Out[32]: [None, None, None, None, None]    
In [33]: idx[2]=slice(None)

In [34]: (M*N[tuple(idx)]).shape
Out[34]: (2, 3, 4, 5, 6)

Or we could reshape N, adding the needed trailing size 1 dimensions:

In [35]: shp = [1]*M.ndim; shp[2]=N.shape[0]
In [36]: (M*N.reshape(shp)).shape
Out[36]: (2, 3, 4, 5, 6)

Or use expand_dims, which uses that reshape under the covers. get the (1,2) from the ndim and desired insert dimension

In [37]: N1=np.expand_dims(N,(1,2)); N1.shape
Out[37]: (4, 1, 1)    
In [38]: (M*N1).shape
Out[38]: (2, 3, 4, 5, 6)

Answered By - hpaulj