Dynamic axis indexing of Numpy ndarray

Question

I want to obtain the 2D slice in a given direction of a 3D array where the direction (or the axis from where the slice is going to be extracted) is given by ano

Answer 1

To index a dimension dynamically, you can use swapaxes, as shown below:

a = np.arange(7 * 8 * 9).reshape((7, 8, 9))

axis = 1
idx = 2

np.swapaxes(a, 0, axis)[idx]

Runtime comparison

Natural method (non dynamic) :

%timeit a[:, idx, :]
300 ns ± 1.58 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

swapaxes:

%timeit np.swapaxes(a, 0, axis)[idx]
752 ns ± 4.54 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

Index with list comprehension:

%timeit a[[idx if i==axis else slice(None) for i in range(a.ndim)]]