Indexing multi-dimensional array with tuple of indices from an indexing array - 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?

Divakar · Accepted Answer · 2019-07-08 16:31:18Z

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)

edited Jul 8, 2019 at 16:31

answered Jul 8, 2019 at 16:04

Divakar

222k19 gold badges273 silver badges374 bronze badges

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

6 Comments

Mad Physicist Over a year ago

Brilliant as usual

rboz22 Over a year ago

Thank you! Why does tuple() work on a 2-D array but not a 3-D array? I assumed it only worked for 1-D.

Divakar Over a year ago

@rboz22 It does work on 3D array too as shown in the b[tuple(a)] method. After the tuple conversion, the packing is elementwise between those tuples to form an indexer and select one element each off b. This is also mentioned in the post.

rboz22 Over a year ago

Sorry, I posted the comment before the edit was made. For n-D arrays, along which dimension does tuple() convert values to indices? tuple(a) seems to yield very large values instead of 0<=i<20.

Divakar Over a year ago

@rboz22 tuple(a) would only collect elements off a into tuples. If elements in a are in 0<element_in_a<20, then tuple(a) would also be restricted to that 0-to-20 range.

|

Collectives™ on Stack Overflow

Indexing multi-dimensional array with tuple of indices from an indexing array - NumPy / Python

1 Answer 1

6 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

6 Comments

Your Answer

Sign up or log in

Post as a guest

Related