Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

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

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!

like image 722
Joe Avatar asked Dec 06 '25 09:12

Joe


1 Answers

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])
like image 194
Andy Hayden Avatar answered Dec 07 '25 22:12

Andy Hayden



Donate For Us

If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!