ati February 2016

What is the nicest way to weight planes in a numpy array?

I have the following code in which w is a 1D numpy array of compatible dimension, and M is a 4D array,

i = 0
for weight in w:
    M[:, :, i, :] *= weight
    i += 1

Is there a nicer way to achieve the same effect?


Divakar February 2016

You are scaling M along axis=2 with the elements from w, which is a 1D array. So, you need to extend w to a 2D array with np.newaxis/None, which will align the axes between extended version of w with M. Then, perform element-wise multiplication between these two arrays to bring in broadcasting for a vectorized solution, like so -

M *= w[:,None]

If axis=2 of M has a length that is more than the number of elements in w, you need to select a range along axis=2 in M before multiplying, like so -

M[...,np.arange(w.size),:] *= w[:,None]

Wolfgang February 2016

This answer is based on my trying to understand @Divakar's answer. What helped me understand what is going on is writing their suggestion of

M *= w[:,None]


M *= w[None,None,:,None]

where now the dimensions of M and the extended w are visibly the same. Of course, @Divakar's version is shorter and so more elegant, but less intuitive.

So, a full working example would be:

import numpy as np
M = np.ones((1,4,3,2))
w = np.arange(3)
M *= w[None,None,:,None]
print M

Post Status

Asked in February 2016
Viewed 3,179 times
Voted 14
Answered 2 times


Leave an answer