numpynumpy-indexing

Get 2D elements from 3D numpy array, corresponding to 2D indices

I have a 3D array that could be interpreted as a 2D matrix of positions, where each position is a 2D array of coordinates [x,y].

I then have a list of 2D indices, each one indicating a position in the matrix in terms of [row, column]. I would like to obtain the positions from the matrix corresponding to all these indices.

What I am doing is:

import numpy as np

input_matrix = np.array(
    [[[0.0, 1.5], [3.0, 3.0]], [[7.0, 5.2], [6.0, 7.0]]]
)

indices = np.array([[1, 0], [1, 1]])

selected_elements = np.array([input_matrix[tuple(idx)] for idx in indices])

So for example the 2D element corresponding to the 2D index [1, 0] would be [7.0, 5.2] and so on.

My code works, but I was wondering if there is a better and way, for example using numpy entirely (e.g. without having to use list comprehension in the case of multiple 2D indices).

I tried to use the numpy take but it does not seem to produce the wanted results.

Solution

  • You can use:

    input_matrix[tuple(indices.T)]

    Or, as suggested in comments:

    input_matrix[indices[:,0], indices[:,1]]

    Output:

    array([[7. , 5.2],
       [6. , 7. ]])