Index a numpy array using a tuple (or list) when len < ndim?

Question

I have a 3d numpy array, eg:

>>> A = np.arange(24).reshape(2,3,4)

I want to take a 1d slice along axis 0 based on a pair of coordinates for axes 1 and 2:

>>> h = 1
>>> l = 2
>>> A[:,h,l]
array([ 6, 18])

So far so good. But what if my coordinate pair is stored as a tuple or a list, rather than two integers? I've experimented with a few obvious options, to no avail:

>>> coords = (1,2)
>>> A[coords]
array([20, 21, 22, 23])
>>> A[:,coords]
array([[[ 4,  5,  6,  7],
        [ 8,  9, 10, 11]],

       [[16, 17, 18, 19],
        [20, 21, 22, 23]]])
>>> A[...,coords]
array([[[ 1,  2],
        [ 5,  6],
        [ 9, 10]],

       [[13, 14],
        [17, 18],
        [21, 22]]])

I've googled around on this and not found anything, but it's entirely possible that I'm not searching with the appropriate jargon. So, apologies if this is an overly simplistic question!

You'll need to unpack your coords into two separate arrays, one for each of the last two axes. E.g. h, l = coords. (Note that coords could be two-dimensional, in this case.) Or am I misunderstanding the question? — Joe Kington
– Joe Kington, Commented Dec 18, 2014 at 21:55

Andy Hayden · Accepted Answer · 2014-12-18 21:56:47Z

5

You can construct the slice tuple directly, with something like:

In [11]: A[(slice(None),) + coords]
Out[11]: array([ 6, 18])

This is because calling A[:, 1, 2] is equivalent / calls:

In [12]: A.__getitem__((slice(None, None, None), 1, 2))
Out[12]: array([ 6, 18])

answered Dec 18, 2014 at 21:56

Andy Hayden

378k110 gold badges640 silver badges546 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Joe Kington Over a year ago

Just on a side note to the OP: it does strictly have to be a tuple of slices, not just any iterable. Numpy will interpret an array or list passed in in an entirely different manner. Nice explanation, as well!

hpaulj Over a year ago

The caution about using tuples is good. However, ind=[slice(None)]*3; ind[1:]=coords; A[ind] works just as well. There's a note in docs.scipy.org/doc/numpy/reference/arrays.indexing.html about this retaining backward compatibility with old numeric package (for a list containing a slice).

Collectives™ on Stack Overflow

Index a numpy array using a tuple (or list) when len < ndim?

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related