CODEWITHSUNDEEP
numpyindexing
Imanol Luengo
Imanol Luengo

Reputation: 15889

Dynamic axis indexing of Numpy ndarray

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 another variable.

Assuming idx the index of the 2D slice in a 3D array, and direction the axis in which obtain that 2D slice, the initial approach would be:

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

I'm pretty sure there must be a way of doing this without doing conditionals, or at least, not in raw python. Does numpy have a shortcut for this?

The better solution I've found so far (for doing it dynamically), relies in 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 there is any better/easier/faster function for this, since for 3D the transpose code almost looks more awful than the 3 if/elif. It generalizes better for ND and the larger the N the more beautiful the code gets in comparison, but for 3D is quite the same.

Upvotes: 13

Views: 2314

Answers (3)

AturSams
AturSams

Reputation: 7952

This is python. You could simply use eval() like this:

def get_by_axis(a, idx, axis):
    indexing_list = a.ndim*[':']
    indexing_list[axis] = str(idx) 
    expression = f"a[{', '.join(indexing_list)}]"
    return eval(expression)

Obviously, in which case you do not accept input from untrusted users.

Upvotes: -1

ndou
ndou

Reputation: 1178

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)]]

Upvotes: -1

hpaulj
hpaulj

Reputation: 231385

Transpose is cheap (timewise). There are numpy functions that use it to move the operational axis (or axes) to a known location - usually the front or end of the shape list. tensordot is one that comes to mind.

Other functions construct an indexing tuple. They may start with a list or array for ease of manipulation, and then turn it into a tuple for application. For example

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

np.apply_along_axis does something like that. It's instructive to look at the code for functions like this.

I imagine the writers of the numpy functions worried most about whether it works robustly, secondly about speed, and lastly whether it looks pretty. You can bury all kinds of ugly code in a function!.

tensordot ends with

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 calculated newaxes_.. and newshape....

apply_along_axis constructs a (0...,:,0...) index tuple

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

Upvotes: 11

Related Questions