ENH: stats._xp_mean, an array API compatible mean with weights and nan_policy
#20743
Conversation
mdhaber
added
scipy.stats
enhancement
| (xp_mean_1samp, tuple(), dict(), 1, 1, False, lambda x: (x,)), | ||
| (xp_mean_2samp, tuple(), dict(), 2, 1, True, lambda x: (x,)), |
There was a problem hiding this comment.
Most scipy.stats functions use the _axis_nan_policy decorator to implement nan_policy, keepdims, and tuple axis. I've implemented all these features natively for improved performance (e.g. nan_policy='omit' would otherwise loop over each slice), and the function still passes all the tests, which are quite stringent. So if you don't want to write tests with hypothesis, I'm still pretty comfortable with this.
scipy/_lib/_array_api.py
Outdated
|
|
||
| if weights is not None and x.shape != weights.shape: | ||
| try: | ||
| x, weights = xp.broadcast_arrays(x, weights) |
There was a problem hiding this comment.
A few thoughts about broadcasting:
- Technically
x = [1, 2, 3]is broadcastable withweights = [2], and it can be interpreted as giving all observations a weight of2. - Technically,
x = [1]is broadcastable withweights = [1, 2, 3]: now we havexbeing broadcast to the shape ofweightsrather than the (more natural) other way around. - Technically
x = []is broadcastable withweights = [1]:weightsgets broadcasted to shape(0,), and the weighted mean is NaN.
It's clearly simpler to just accept these sorts of things, but since they're not useful, one could argue that we shouldn't. I'd propose that we just accept them, but if there are strong opinions about not accepting them, LMK.
scipy/_lib/_array_api.py
Outdated
| @@ -475,3 +476,155 @@ def xp_sign(x, xp=None): | |||
| sign = xp.where(x < 0, -one, sign) | |||
| sign = xp.where(x == 0, 0*one, sign) | |||
| return sign | |||
|
|
|||
|
|
|||
| def xp_add_reduced_axes(res, axis, initial_shape, *, xp=None): | |||
There was a problem hiding this comment.
Could you add a note on why this is needed? Is it temporary, why can't xp.add not be used, etc.?
There was a problem hiding this comment.
Type annotations and consistency with other functions in this file would be useful too (at least if you expect this function to stay around for a while).
res should preferably be positional-only.
There was a problem hiding this comment.
Perhaps a better name would have been xp_replace_reduced_axes or xp_keepdims: it adds back axes that have been reduced away. However, when there are other comments to respond to, I'll just move the logic back into xp_mean, since I'm not sure if it will be used elsewhere. It can be factored out again as needed. Although the comment wasn't about xp_mean, I can make the first argument of xp_mean positional-only.
|
Why was |
|
|
Thanks for the link @lucascolley . To align with the naming convention of |
lucascolley
changed the title
mean with weights and nan_policyxp_mean, an array API compatible mean with weights and nan_policy
mdhaber
changed the title
xp_mean, an array API compatible mean with weights and nan_policystats._xp_mean, an array API compatible mean with weights and nan_policy
Reference issue
Toward gh-20544
What does this implement/fix?
This function adds
_xp_mean, an array-API compatible function which combines the features ofnp.mean,np.average, andnp.nanmeanin interface that fits withscipy.stats. This will be needed for making functions likepmean,hmean, andgmeanarray-API compatible.Additional information
Potential reviewers: would you be willing to write some unit tests withhypothesis? For such a fundamental function, it's particularly important that it works flawlessly!If it doesn't sound too crazy, I'd suggest that this and similarvarandstdfunctions be added publicly toscipy.statsbecause they provide functionality that does not exist with the array API (e.g.weights, which has been explicitly rejected, andnan_policy, which has not been standardized and may not follow SciPy's convention). Even considering NumPy alone, it would be useful to have a single function that has all the functionality ofmean,average, andnanmeanin an interface consistent with the rest ofscipy.stats.Not pursuing these things right now. Let's just get this in so we can finish the other mean functions.