使用索引数组中的索引元组索引多维数组 - NumPy / Python

Question

I have a 3-D numpy array a with dimensions (6,m,n) . I also have a 6-D boolean numpy array b with dimensions (20,20,20,20,20,20) that effectively works as a mask.

I would like to use the 6 values at each location (m,n) in the first array to retrieve the corresponding value in the second array. Effectively, I will compress the 3D int array into a 2D boolean array. I thought the solution would be using np.where , but I don't think it can deal with using values as indices.

The naive implementation for this will be something like:

for i in range(m):
    for j in range(n):
         new_arr[i,j]=b[tuple(a[:,i,j])]

Is there any way to implement this without using a loop?

Answer 1

Approach #1

Reshape a to 2D keeping the first axis length as the same. Convert each thus 2D-flattened-block to a tuple and then index into b . This tuple-conversion leads to a packing of each elements along the first axis as an indexer to select an element each off b . Finally a reshaping is needed to get a 2D output. Hence, the implementation would look something like this -

b[tuple(a.reshape(6,-1))].reshape(m,n)

Or, skip all that reshaping mess and simply do -

b[tuple(a)]

This does the same indexer creation and solves the problem.

Approach #2

Alternatively, we can also compute the flattened indices and then index into flattened b with those and extract relevant boolean values off it -

b.ravel()[np.ravel_multi_index(a,b.shape)]

Timings on a large dataset -

In [89]: np.random.seed(0)
    ...: m,n = 500,500
    ...: b = np.random.rand(20,20,20,20,20,20)>0.5
    ...: a = np.random.randint(0,20,(6,m,n))

In [90]: %timeit b[tuple(a)]
14.6 ms ± 184 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [91]: %timeit b.ravel()[np.ravel_multi_index(a,b.shape)]
7.35 ms ± 136 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)