How can I extract values from a 3D numpy array using a 2D numpy array with indices?

I have two numpy arrays of the following shapes, and I need to extract values from a 3D array using stored indices from a 2D one:

• vals = (65, 65, 3500) This contains 3500 values that were calculated in a 65×65 grid, from which I need to extract values
• indices = (65, 65) This contains the indices of the minimum value of each of the 3500 of the previous array, of which I need to extract the values

The resulting array should again be (65, 65).

For example, if I have a value 1367 stored in indices[0,0], then I need vals[0,0,1367]. This needs to be repeated all the way to indices[64,64].

How can I do this sort of filter? I tried np.choose with reshaping vals to (3500, 65, 65), but that is limited to 32 columns and therefore crashed. Alternatively, how can I get the (65, 65) array of the minima of each set of 3500 values? Thanks

Answer

import numpy as np
vals.reshape(-1, 3500)[np.arange(65*65), indices.ravel()].reshape(65, 65)

This simplifies the operation by only doing selection by row. On my PC this is also slightly faster than the meshgrid method.

36 µs ± 239 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

compared to

44.6 µs ± 693 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)