2

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
Timothy W. Hilton
  • 431
  • 2
  • 11
  • `data.swapaxes(0, 1)[idx]`? – cs95 Aug 03 '17 at 16:43
  • https://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 Aug 03 '17 at 18:05

1 Answers1

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?

hpaulj
  • 221,503
  • 14
  • 230
  • 353