Spencer Hill Spencer Hill - 3 months ago 25
Python Question

Numpy averaging with multi-dimensional weights along an axis

I have a numpy array,

a
,
a.shape=(48,90,144)
. I want to take the weighted average of
a
along the first axis using the weights in array
b
,
b.shape=(90,144)
. So the output should be a numpy array of shape
(48,)
.

I know this can be done with a list comprehension:

np.array([np.average(a[i], weights=b) for i in range(48)])


But I'd like to avoid having to convert from a list back to a numpy array.

Can anyone help? I'm sure this is possible using numpy functions and slicing, but I'm stuck. Thanks!

Answer

In a single line:

np.average(a.reshape(48, -1), weights=b.ravel()), axis=1)

You can test it with:

a = np.random.rand(48, 90, 144)
b = np.random.rand(90,144)
np.testing.assert_almost_equal(np.average(a.reshape(48, -1),
                                          weights=b.ravel(), axis=1),
                               np.array([np.average(a[i],
                                                    weights=b) for i in range(48)]))
Comments