numpy ndarray advanced indexing

Question

I have ndarray of 3 dimension. How do I select index 0 and 1 from first axis while selecting index 0 and 3 from second axis and index 1 from third axis?

I tried to use index [(0,1), (1, 3), 1], which produces a result completely different than what I thought it would produce.

So two questions here. What does [(0,1), (1, 3), 1] do? And how to correctly create an index that solve my original question.

a = np.arange(30).reshape(3, 5, 2)
array([[[ 0,  1],
        [ 2,  3],
        [ 4,  5],
        [ 6,  7],
        [ 8,  9]],

       [[10, 11],
        [12, 13],
        [14, 15],
        [16, 17],
        [18, 19]],

       [[20, 21],
        [22, 23],
        [24, 25],
        [26, 27],
        [28, 29]]])

a[0, (1, 3), 1]  # produces array([3, 7])
a[(0,1), (1, 3), 1] # produces array([ 3, 17])

```

...and what's your expected output?

Ashwini Chaudhary
– Ashwini Chaudhary

2015-06-03 02:35:37 +00:00
Commented Jun 3, 2015 at 2:35 — Ashwini Chaudhary
– Ashwini Chaudhary, Commented Jun 3, 2015 at 2:35

user2357112 · Accepted Answer · 2015-06-03 06:00:20Z

1

When you index the way you're doing it, NumPy doesn't interpret it as selecting those indices of each dimension. Instead, NumPy broadcasts the arguments against each other:

a[(0,1), (1, 3), 1] -> a[array([0, 1]), array([1, 3]), array([1, 1])]

and then creates a result array where a[i, j, k][x] == a[i[x], j[x], k[x]].

To get the behavior you're looking for, you need to reshape the arguments you're passing, so that broadcasting them against each other produces an array of shape (2, 2) instead of shape (2,). This means the first argument needs shape (2, 1), the second argument needs to have shape (1, 2) or (2,), and the third argument's shape is fine. numpy.ix_ can make this easier, but it doesn't support scalar arguments. a[np.ix_([0, 1], [1, 3], [1])] does what you would have expected a[[0, 1], [1, 3], [1]] to do, but to get the shape you would have expected from a[[0, 1], [1, 3], 1], your options are messier:

>>> a[np.ix_([0, 1], [1, 3], [1])]
array([[[ 3],
        [ 7]],

       [[13],
        [17]]])
>>> a[np.ix_([0, 1], [1, 3]) + (1,)]
array([[ 3,  7],
       [13, 17]])
>>> a[np.ix_([0, 1], [1, 3], [1])][:, :, 0]
array([[ 3,  7],
       [13, 17]])

edited Jun 3, 2015 at 6:00

answered Jun 3, 2015 at 2:35

user2357112

286k32 gold badges490 silver badges569 bronze badges

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

1 Comment

Negative Zero Over a year ago

Thanks, that explains. I just read up about the numpy broadcasting. ix_ function looks good and seems to satisfies what I need.

Collectives™ on Stack Overflow

numpy ndarray advanced indexing

1 Answer 1

1 Comment

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Linked

Related