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