Numpy ndarray dynamic axis indexing

Question

I want to get a 2D slice in the given direction of a 3D array, where direction (or the axis from which the slice will be extracted) is set by another variable.

Assuming idx index of the two-dimensional slice in the 3D array and the direction axis in which this 2D slice is obtained, the initial approach:

 if direction == 0: return A[idx, :, :] elif direction == 1: return A[:, idx, :] else: return A[:, :, idx] 

I am sure there must be a way to do this without making conventions, or at least not in raw python. Does numpy have a shortcut for this?

The best solution I've found so far (for its dynamic) relies on the transpose operator:

 # for 3 dimensions [0,1,2] and direction == 1 --> [1, 0, 2] tr = [direction] + range(A.ndim) del tr[direction+1] return np.transpose(A, tr)[idx] 

But I wonder if the function is better / easier / faster for this, since for 3D the transpose code looks almost terrible than 3 if / elif. It generalizes ND better, and the more N, the more beautiful the code is in comparison, but for 3D it is exactly the same.

Answer 1

Transposing is cheap (in time). There are numpy functions that use it to move a working axis (or axes) to a known location — usually the front or the end of a list of shapes. tensordot is one that comes to mind.

Other functions build an indexing tuple. They can start with a list or array for easy manipulation, and then turn it into a tuple for the application. for instance

 I = [slice(None)]*A.ndim I[axis] = idx A[tuple(I)] 

np.apply_along_axis does something similar. It is instructive to look at the code for such functions.

I believe that the authors of numpy functions were most worried about whether it works decisively, and secondly, about speed and, finally, whether it looks beautiful. You can bury all kinds of ugly code in a function!

---

tensordot ends on

 at = a.transpose(newaxes_a).reshape(newshape_a) bt = b.transpose(newaxes_b).reshape(newshape_b) res = dot(at, bt) return res.reshape(olda + oldb) 

where the previous code is calculated by newaxes_.. and newshape...

apply_along_axis builds an index tuple (0...,:,0...)

 i = zeros(nd, 'O') i[axis] = slice(None, None) i.put(indlist, ind) ....arr[tuple(i.tolist())] 

Answer 2

To dynamically index a dimension, 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 (not 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) 

List Pointer:

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