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 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.
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 ?
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))
Collected from the Internet
Please contact [email protected] to delete if infringement.
Comments