-
-
Notifications
You must be signed in to change notification settings - Fork 10.8k
MAINT: np.pad: Add helper functions for producing slices along axes #11012
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
@eric-wieser What is the current status of this? |
Needs a base-switch from master and back again now that the other PR is merged, which will hide the first merged commit from the diff |
numpy/lib/arraypad.py
Outdated
@@ -771,8 +764,7 @@ def _pad_ref(arr, pad_amt, method, axis=-1): | |||
ref_chunk2 = arr[ref_slice][rev_idx] | |||
|
|||
if 'odd' in method: | |||
edge_slice2 = tuple(slice(None) if i != axis else slice(x - 1, x) | |||
for (i, x) in enumerate(arr.shape)) | |||
edge_slice2 = _end_slice(arr.shape, 1, axis=axis) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The alternative I was considering was _axis_slice(arr.shape, lambda n: slice(n-1, n), axis=axis)
, which covers a few more uses, but I wasn't sure if that was clearer.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Not clearer as written, but I think it would be clearer if one were to just pass in the required slice itself. This is possible, since, as I noted above, it is not necessary to know the actual dimension; here, one would just pass in slice(-1, None)
.
numpy/lib/arraypad.py
Outdated
for (i, x) in enumerate(shape)) | ||
|
||
|
||
def _end_slice(shape, n, axis): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
One doesn't actually need the shape
for either _begin_slice
(obvious, it is not used), or _end_slice
: just do
def _end_slice(ndim, n, axis):
return tuple(slice(None) if i != axis else slice(-n, None) for i in range(ndim))
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
-n
fails when n == 0
. I'd like to remove some of the zero special-casing, if possible
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actually, thinking more, we do not need ndim
either, since you only need a tuple up to axis
- all the later dimensions are automatically taken to be slice(None)
. So, it could be:
def _slice_at_axis(axis, slice):
return (slice(None),) * axis + (slice,)
(This is substantially faster than the loop, at least for a randomly selected axis=3
.)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Perhaps pass in slice(arr.shape[axis] - padding, None)
in that case - at least it is immediately obvious what the code does.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I was trying to avoid writing axis
twice, hence the lambda function above.
This is substantially faster than the loop
I suspect this would be true if I kept my existing function names but changed their implementations too (I deliberately didn't do that in this PR, because it makes the diff clearer)
2d52077
to
7dd7938
Compare
Updated to build my helper functions out of the simpler helper function you propose (which is useful for a few other pieces of cleanup. Also renamed from |
e25e01b
to
944f1b4
Compare
Construct a slice tuple the length of shape, with sl at the specified axis | ||
""" | ||
slice_tup = (slice(None),) | ||
return slice_tup * axis + (sl,) + slice_tup * (len(shape) - axis - 1) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, the last half is redundant - but it produces the same slices as we did before, and I see little reason to change that.
This makes `_slice_first` almost a factor of two faster
944f1b4
to
5608636
Compare
if num is not None: | ||
mean_slice = tuple( | ||
slice(None) if i != axis else slice(end, end - num, -1) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Note that this slice is the reverse of what it was before - but min
, max
, mean
, and med
all are order-invariant anyway
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'm a bit surprised by the double slicing, but that was there already and I think this is all OK now, so approved! Maybe squash the two commits, though?
I think it's clearer as is with two commits - extracting helper functions from the existing code, then rewriting those helpers to be more efficient |
OK, that's fine too. Will merge... |
Thanks for the feedback! |
Follows on from #11011