PR to allow weights that align with specified axis in np.average()
https://github.com/numpy/numpy/pull/25541/ np.average() and weightsSummary This PR intends to address the limitation that the weights supplied as an optional argument in np.average() can only be the same shape as the input array, or onedimensional. With these changes, weights that have the same shape as the input array over the specified axis can be supplied. The PR also makes some minor improvements to the sanity checks on the weights. We have implemented the extension for both np.average() and np.ma.average(), and added tests for the new functionality (all preexisting tests pass). Details Some functionalities of the numpy.average() function can be improved, and the treatment of weights in particular. Currently, np.average() is restricted to weights that are either 1. the same shape as the input array, or 2. onedimensional. This restriction is not necessary and is disadvantageous when averaging over subdimensions the input array via supplying the axis argument. If it's possible to average over subdimensions of the input array via axis, then it is useful to be able to supply weights in accordance, checking that they have the required dimensionality and shape. Consider, for example, y = np.arange(12).reshape(2, 2, 3)w = np.array([0., 0., 1., .5, .5, 0., 0., .5, .5, 1., 0., 0.]).reshape(2, 2, 3) both the input array y and the weights w have shape (2, 2, 3). If I want to average over the whole array, that is OK: np.average(y, weights=w) 5.5 or over a onedimensional subset (axis=0): np.average(y, axis=0 weights=w[:, 0, 1]) array([[ 6., 7., 8.], [ 9., 10., 11.]]) If, however, I want to average a subdimensions of y that are not onedimensional, and supply weights with the same shape as the subdimensions, that is not possible: np.average(y, axis=(0, 1) weights=w[:, :, 1]) TypeError: 1D weights expected when shapes of a and weights differ. Of course, I can add an extra dimension to my weights to make them the same shape of y and get around this limitation, however, depending on the shape ofa in the dimension that is *not* being averaged over, that might be very expensive in terms of memory and perhaps prohibitively so (it was for me in one case). This PR addresses this limitation. Moreover, some of the checks on the weights could be improved, for example with the previous implementation if I supplied 1d weights but a tuple for the axis: np.average(y, axis=(0, 1), weights=w[0, 0, :]) TypeError: tuple indices must be integers or slices, not tuple This issue is not caught by one of the sanity checks and the error message does not clearly tell the user what has gone wrong. This PR improves that as well. The main idea was to generalise the function np.average(), but I ended up doing the same for np.ma.average() too, for completeness, even though I have relatively little experience with masked arrays. The implementations have been tested in lib/tests/test_function_base.py and ma/tests/test_extras.py, respectively. The two sets of tests mirror one another. All preexisting tests pass. I'd appreciate if people could kindly take a look at this PR and let know any feedback they have. Thank you James
participants (1)

James Hamp