Numpy: "Adding" index and colon

Question

I have

>>> foo = np.zeros((3,3,3))
>>> foo[1,2,1] = 1
>>> idx = 1,2

I would like to get the equivalent of

>>> foo[1,2,:]
array([ 0.,  1.,  0.])

using idx (to iterate through idx) . Both approaches that I tried didn't work out:

>>> foo[idx, :]
array([[[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  1.,  0.]],

       [[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  0.]]])
>>> foo[((idx,)+(slice(None),))]
array([[[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  1.,  0.]],

       [[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  0.]]])

hpaulj · Accepted Answer · 2016-01-17 17:45:14Z

5

Use foo[idx]. Read the following to see how I arrived at this.

foo[1,2,:] is effectively:

In [379]: foo[(1,2,slice(None))]
Out[379]: array([ 0.,  1.,  0.])

The Python interpreter converts the 1,2,: to this tuple, and passes it to the foo.__getitem__ method.

So we just have to find the right way of constructing the tuple. One is:

In [380]: tuple(idx)+(slice(None),)
Out[380]: (1, 2, slice(None, None, None))

resulting in:

In [381]: foo[tuple(idx)+(slice(None),)]
Out[381]: array([ 0.,  1.,  0.])

Actually I don't need the tuple call, idx is already a tuple

In [386]: idx
Out[386]: (1, 2)
In [387]: idx+(slice(None),)
Out[387]: (1, 2, slice(None, None, None))

I would need tuple([1,2]) if idx was initialed as a list.

Double actually, the solution is even simpler. Since we are using : for the last dimension, we can omit it. And since idx is tuple, it already indexes the first 2 dimensions.

In [394]: foo[idx]
Out[394]: array([ 0.,  1.,  0.])

This pair of calls may lend some clarity:

In [396]: foo[(1,2)]
Out[396]: array([ 0.,  1.,  0.])
In [397]: foo[[1,2]]
Out[397]: 
array([[[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  1.,  0.]],

       [[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  0.]]])

edited Jan 17, 2016 at 17:45

answered Jan 17, 2016 at 17:39

hpaulj

233k14 gold badges260 silver badges392 bronze badges

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

1 Comment

FooBar Over a year ago

Is there a generalization of this, say for 1,2, 0:2?

Ignacio Vazquez-Abrams · Accepted Answer · 2016-01-17 12:21:04Z

2

>>> foo[idx + (slice(None),)]
array([ 0.,  1.,  0.])

answered Jan 17, 2016 at 12:21

Ignacio Vazquez-Abrams

804k160 gold badges1.4k silver badges1.4k bronze badges

1 Comment

hpaulj Over a year ago

A bit of explanation about tuples and slice objects might help the poster.

Collectives™ on Stack Overflow

Numpy: "Adding" index and colon

2 Answers 2

1 Comment

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related