I have a numpy ndarray arr
and also indices
, a list of indices specifying a particular entry. For concreteness let's take:
arr = np.arange(2*3*4).reshape((2,3,4))
indices= [1,0,3]
I have code to take 1d slices through arr
observing all but one index n
:
arr[:, indices[1], indices[2]] # n = 0
arr[indices[0], :, indices[2]] # n = 1
arr[indices[0], indices[1], :] # n = 2
I would like to change my code to loop over n
and support an arr
of arbitrary dimension.
I've had a look at the indexing routines entry in the documentation and found information about slice()
and np.s_()
. I was able to hack together something that works like what I want:
def make_custom_slice(n, indices):
s = list()
for i, idx in enumerate(indices):
if i == n:
s.append(slice(None))
else:
s.append(slice(idx, idx+1))
return tuple(s)
for n in range(arr.ndim):
np.squeeze(arr[make_custom_slice(n, indices)])
Where np.squeeze
is used to remove the axes of length 1. Without this, the array this produced has shape (arr.shape[n],1,1,...)
rather than (arr.shape[n],)
.
Is there a more idiomatic way to accomplish this task?
Some improvements to the solution above (there may still be a one-liner or a more performant solution):
def make_custom_slice(n, indices):
s = indices.copy()
s[n] = slice(None)
return tuple(s)
for n in range(arr.ndim):
print(arr[make_custom_slice(n, indices)])
An integer value idx
can be used to replace the slice object slice(idx, idx+1)
. Because most indices are copied over directly, start with a copy of indices rather than building the list from scratch.
When built in this way, the result of arr[make_custom_slice(n, indices)
has the expected dimension and np.squeeze
is unnecessary.
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.