Indexing array with array on numpy

Question

It is similar to some questions around SO, but I don't quite understand the trick to get what I want.

I have two arrays,
arr of shape (x, y, z)
indexes of shape (x, y) which hold indexes of interest for z.

For each value of indexes I want to get the actual value in arr where:

arr.x == indexes.x  
arr.y == indexes.y  
arr.z == indexes[x,y]

This would give an array of shape(x,y) similar to indexes' shape.

For example:

arr = np.arange(99)
arr = arr.reshape(3,3,11)
indexes = np.asarray([
[0,2,2],
[1,2,3],
[3,2,10]])
# indexes.shape == (3,3)

# Example for the first element to be computed
first_element = arr[0,0,indexes[0,0]]

With the above indexes, the expected arrays would look like:

expected_result = np.asarray([
[0,13,24],
[34,46,58],
[69,79,98]])

I tried elements = np.take(arr, indexes, axis=z) but it gives an array of shape (x, y, x, y)
I also tried things like elements = arr[indexes, indexes,:] but I don't get what I wish.
I saw a few answers involving transposing indexes and transforming it into tuples but I don't understand how it would help.

Note: I'm a bit new to numpy so I don't fully understand indexing yet.
How would you solve this numpy style ?

"but I don't get what I wish": It is not clear what you wish exactly. Can you give an example with some dummy data? E.g. if `arr = np.arange(24).reshape(2,3,4)`, then what would `indexes` look like? And what is the expected result then? — Tawy, Sep 01 '21 at 13:56
What would be the resulting array in this case, you have only given the first element of the output. — Ivan, Sep 01 '21 at 14:15

bb1 · Accepted Answer · 2021-09-01T14:55:52.813

This can be done using np.take_along_axis

import numpy as np

#sample data
np.random.seed(0)
arr = np.arange(3*4*2).reshape(3, 4, 2) # 3d array
idx = np.random.randint(0, 2, (3, 4))   # array of indices

out = np.squeeze(np.take_along_axis(arr, idx[..., np.newaxis], axis=-1))

In this code, the array of indices gets added one more axis, so it can be broadcasted to the shape of the array arr from which we are making the selection. Then, since the return value of np.take_along_axis has the same shape as the array of indices, we need to remove this extra dimension using np.squeeze.

Another option is to use np.choose, but in this case the axis along which you are making selections must be moved to be the first axis of the array:

out = np.choose(idx, np.moveaxis(arr, -1, 0))

Thanks, it does wonders ! I'm accepting this answer because it should work for n dimensional array not only 3d (even though it was not in my question) and there is no need to reshape. — Abel, Sep 01 '21 at 14:44

score 1 · Answer 2 · answered Sep 01 '21 at 14:20

1

The solution here should work for you: Indexing 3d numpy array with 2d array

Adapted to your code:

ax_0 = np.arange(arr.shape[0])[:,None]
ax_1 = np.arange(arr.shape[1])[None,:]

new_array = arr[ax_0, ax_1, indexes]

answered Sep 01 '21 at 14:20

mdgrogan

169
4

Ivan · Answer 3 · 2021-09-01T14:31:06.057

You can perform such an operation with np.take_along_axis, the operation can only be applied along one dimension so you will need to reshape your input and indices.

The operation you are looking to perform is:

out[i, j] = arr[i, j, indices[i, j]]

However, we are forced to reshape both arr and indices, i.e. map (i, j) to k, such that we can apply np.take_along_axis. The following operation will take place:

out[k] = arr[k, indices[k]] # indexing along axis=1

The actual usage here comes down to:

>>> put = np.take_along_axis(arr.reshape(9, 11), indices.reshape(9, 1), axis=1)
array([[ 0],
       [13],
       [24],
       [34],
       [46],
       [58],
       [69],
       [79],
       [91]])

Then reshape back to the shape of indices:

>>> put.reshape(indices.shape)
array([[ 0, 13, 24],
       [34, 46, 58],
       [69, 79, 91]])

Indexing array with array on numpy

3 Answers3