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

Question

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

Asked By: user243896

||

Source

Answer 1

this can be done using a meshgrid

import numpy as np

vals = np.random.rand(65,65,300)
indices  = np.random.randint(0,300,[65,65])
mesh_grid = np.meshgrid(np.arange(indices.shape[1]),np.arange(indices.shape[0]),sparse=True)
output = vals[mesh_grid[1],mesh_grid[0],indices]

print(output.shape)

(65, 65)

Answered By: Ahmed AEK

Answer 2

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)

Answered By: Z Li

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

Question:

Answers: