indexing a numpy array with indices from another array

Question

I have a numpy array "data" with dimensions [t, z, x, y]. The dimensions represent time (t) and three spatial dimensions (x, y, z). I have a separate array "idx" of indices with dimensions [t, x, y] describing vertical coordinates in data: each value in idx describes a single vertical level in data.

I want to pull out the values from data indexed by idx. I've done it successfully using loops (below). I've read several SO threads and numpy's indexing docs but I haven't been able to make it more pythonic/vectorized.

Is there an easy way I'm just not getting quite right? Or maybe loops are a clearer way to do this anyway...

import numpy as np

dim = (4, 4, 4, 4)  # dimensions time, Z, X, Y

data = np.random.randint(0, 10, dim)
idx = np.random.randint(0, 3, dim[0:3])

# extract vertical indices in idx from data using loops
foo = np.zeros(dim[0:3])
for this_t in range(dim[0]):
    for this_x in range(dim[2]):
        for this_y in range(dim[3]):
            foo[this_t, this_x, this_y] = data[this_t,
                                               idx[this_t, this_x, this_y],
                                               this_x,
                                               this_y]

# surely there's a better way to do this with fancy indexing
# data[idx] gives me an array with dimensions (4, 4, 4, 4, 4, 4)
# data[idx[:, np.newaxis, ...]] is a little closer
# data[tuple(idx[:, np.newaxis, ...])] doesn't quite get it either
# I tried lots of variations on those ideas but no luck yet

stackoverflow.com/a/44868461/901925 does this sort of assignment with a 3d array. The idea should work with your 4d case. The idea is to get 3 arrays (from ogrid) that let you do data[I, idx, J, K] — hpaulj
– hpaulj, Commented Aug 3, 2017 at 18:05

hpaulj · Accepted Answer · 2017-08-03 19:39:24Z

1

In [7]: I,J,K = np.ogrid[:4,:4,:4]
In [8]: data[I,idx,J,K].shape
Out[8]: (4, 4, 4)
In [9]: np.allclose(foo, data[I,idx,J,K])
Out[9]: True

I,J,K broadcast together to the same shape as idx (4,4,4).

More detail on this kind of indexing at

How to take elements along a given axis, given by their indices?

answered Aug 3, 2017 at 19:39

hpaulj

233k14 gold badges260 silver badges392 bronze badges

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

2 Comments

Timothy W. Hilton Over a year ago

wow, that's brilliant. I had to dig into np.ogrid a little to wrap my head around it, but nicely done. Thanks much.

hpaulj Over a year ago

Print the I,J,K variables. You can recreate those yourself,with expressions like np.arange(4)[None,:,None].

Collectives™ on Stack Overflow

indexing a numpy array with indices from another array

1 Answer 1

2 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related