Indexing numpy array with list of slices

Question

I have a list of slices and use them to index a numpy array.

arr = np.arange(25).reshape(5, 5)
# array([[ 0,  1,  2,  3,  4],
#        [ 5,  6,  7,  8,  9],
#        [10, 11, 12, 13, 14],
#        [15, 16, 17, 18, 19],
#        [20, 21, 22, 23, 24]])

slice_list = list(map(lambda i: slice(i, i+2), [1, 2]))
# [slice(1, 3, None), slice(2, 4, None)]

print(arr[slice_list])
# == arr[1:3, 2:4]
# [[ 7  8]
#  [12 13]]

This works fine but it breaks if I have fewer slices than the number of dimensions of the array I want to index.

arr3d = arr[np.newaxis, :, :]  # dims: [1, 5, 5]
arr3d[:, slice_list]
# IndexError: only integers, slices (`:`), ellipsis (`...`),(`None`) 
#  numpy.newaxis and integer or boolean arrays are valid indices

The following examples work however:

arr3d[:, slice_list[0], slice_list[1]]
arr3d[[slice(None)] + slice_list]
arr3d[:, [[1], [2]], [2, 3]]

Is there a way I can use a list of slices to index an array with more dimensions. I want to do things like:

arr[..., slice_list]
arr[..., slice_list, :]
arr[:, slice_list, :]

without thinking about the dimensions of the array and figuring out how many [slice(None)]*X I have to pad on either side of my slice_list.

javidcf · Accepted Answer · 2020-01-31 10:03:04Z

3

You can do that using tuples of slices and ellipsis objects. Just put all the elements to you want to use for indexing into a tuple and use it as index:

import numpy as np

arr = np.arange(24).reshape(2, 3, 4)
print(arr)
# [[[ 0  1  2  3]
#   [ 4  5  6  7]
#   [ 8  9 10 11]]
# 
#  [[12 13 14 15]
#   [16 17 18 19]
#   [20 21 22 23]]]
slice_tup = tuple(map(lambda i: slice(i, i+2), [1, 2]))
print(slice_tup)
# (slice(1, 3, None), slice(2, 4, None))
print(arr[slice_tup])
# [[[20 21 22 23]]]

# arr[..., slice_list]
print(arr[(Ellipsis, *slice_tup)])
# [[[ 6  7]
#   [10 11]]
# 
#  [[18 19]
#   [22 23]]]

# arr[..., slice_list, :]
print(arr[(Ellipsis, *slice_tup, slice(None))])
# [[[20 21 22 23]]]

# arr[:, slice_list, :]
print(arr[(slice(None), *slice_tup, slice(None))])
# IndexError: too many indices for array

edited Jan 31, 2020 at 10:03

answered Jan 30, 2020 at 16:50

javidcf

59.9k7 gold badges87 silver badges134 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

tenhjo Over a year ago

Ok, but I have to combine the lists/tuples, I can't use arr[..., slice_tuple, :] directly?

javidcf Over a year ago

@scleronomic No, because that would be the same as arr[(Ellipsis, (slice(1, 3, None), slice(2, 4, None)), slice(None))] (that is, a nested tuple). For these cases, it would be nice to have a * unpacking syntax, as with function calls, so you could do something like arr[..., *slice_tup, :], but that is not allowed currently (to be fair is kind of a corner case). But you can do print(arr[(Ellipsis, *slice_tup, slice(None))]), which is probably a bit nicer (actually I think I will update the answer to use that).

Collectives™ on Stack Overflow

Indexing numpy array with list of slices

1 Answer 1

2 Comments

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Related