Skip to content Skip to sidebar Skip to footer

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

Solution 1:

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)]))

Solution 2:

That was the fasted I could come up with:

(a * b).mean(-1).mean(-1) * (b.size / b.sum())

It can be fit for any number of source and result dimensions.

Reshape and 1 x mean did not further speed up:

(a * b).reshape(len(a), -1).mean(-1) * (b.size / b.sum())

Post a Comment for "Numpy Averaging With Multi-dimensional Weights Along An Axis"