numpy take can't index using slice

Question

According to the numpy docs for take it does the same thing as “fancy” indexing (indexing arrays using arrays). However, it can be easier to use if you need elements along a given axis.

However, unlike "fancy" or regular numpy indexing, using slices as indices appears to be not supported:

In [319]: A = np.arange(20).reshape(4, 5)

In [320]: A[..., 1:4]
Out[320]: 
array([[ 1,  2,  3],
       [ 6,  7,  8],
       [11, 12, 13],
       [16, 17, 18]])

In [321]: np.take(A, slice(1, 4), axis=-1)
TypeError: long() argument must be a string or a number, not 'slice'

What is the best way to index an array using slices along an axis only known at runtime?

Answer 1

The most efficient method seems to be A[(slice(None),) * axis + (slice(1, 4),)]:

In [19]: import numpy as np
    ...: x = np.random.normal(0, 1, (50, 50, 50))
    ...: s = slice(10, 20)
    ...: axis = 2
    ...: 
    ...: 

In [20]: timeit np.rollaxis(np.rollaxis(x, axis, 0)[s], 0, axis + 1)
2.32 µs ± 15.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [21]: timeit x.take(np.arange(x.shape[axis])[s], axis)
28.5 µs ± 38.2 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [22]: timeit x[(slice(None),) * axis + (s,)]
321 ns ± 0.341 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)