Slicing a numpy array along a dynamically specified axis
Solution 1:
As it was not mentioned clearly enough (and i was looking for it too):
an equivalent to:
a = my_array[:, :, :, 8]
b = my_array[:, :, :, 2:7]
is:
a = my_array.take(indices=8, axis=3)
b = my_array.take(indices=range(2, 7), axis=3)
Solution 2:
I think one way would be to use slice(None)
:
>>> m = np.arange(2*3*5).reshape((2,3,5))
>>> axis, start, end = 2, 1, 3
>>> target = m[:, :, 1:3]
>>> target
array([[[ 1, 2],
[ 6, 7],
[11, 12]],
[[16, 17],
[21, 22],
[26, 27]]])
>>> slc = [slice(None)] * len(m.shape)
>>> slc[axis] = slice(start, end)
>>> np.allclose(m[slc], target)
True
I have a vague feeling I've used a function for this before, but I can't seem to find it now..