CODEWITHSUNDEEP
pythonnumpy
Abel
Abel

Reputation: 798

Indexing array with array on numpy

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 ?

Upvotes: 2

Views: 604

Answers (3)

bb1
bb1

Reputation: 7873

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))

Upvotes: 3

Ivan
Ivan

Reputation: 40768

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]])

Upvotes: 1

mdgrogan
mdgrogan

Reputation: 179

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]

Upvotes: 1

Related Questions