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 65x65 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

vals[np.arange(65)[:, None], np.arange(65)[None], indices] or vals[(*np.indices(vals.shape[:-1], sparse=True), indices)] — Mechanic Pig
– Mechanic Pig, Commented Oct 13, 2022 at 3:04

Z Li · Accepted Answer · 2022-10-12 22:16:01Z

1

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)

edited Oct 12, 2022 at 22:16

answered Oct 12, 2022 at 22:05

Z Li

4,3381 gold badge7 silver badges21 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

user243896 Over a year ago

What an elegant reply, thank you!

Ahmed AEK · Accepted Answer · 2022-10-13 20:38:05Z

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)

edited Oct 13, 2022 at 20:38

answered Oct 12, 2022 at 22:00

Ahmed AEK

23.2k3 gold badges19 silver badges50 bronze badges

1 Comment

user7711283 Over a year ago

It works now OK after you have fixed it (deleting my first comment to avoid confusion by future visitors).

Collectives™ on Stack Overflow

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

2 Answers 2

1 Comment

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related