7

I have a 3D numpy array data and another array pos of indexes (an index is a numpy array on its own, which makes the latter array a 2D array):

import numpy as np
data = np.arange(8).reshape(2, 2, -1)
#array([[[0, 1],
#    [2, 3]],
#
#  [[4, 5],
#    [6, 7]]])

pos = np.array([[1, 1, 0], [0, 1, 0], [1, 0, 0]])
#array([[1, 1, 0],
#       [0, 1, 0],
#       [1, 0, 0]])

I want to select and/or mutate the elements from data using the indexes from pos. I can do the selection using a for loop or a list comprehension:

[data[tuple(i)] for i in pos]
#[6, 2, 4]
data[[i for i in pos.T]]
#array([6, 2, 4])

But this does not seem to be a numpy way. Is there a vectorized numpy solution to this problem?

cs95
  • 379,657
  • 97
  • 704
  • 746
DYZ
  • 55,249
  • 10
  • 64
  • 93

1 Answers1

3

You can split pos into 3 separate arrays and index, like so—

>>> i, j, k = pos.T
>>> data[i, j, k]
array([6, 2, 4])

Here, the number of columns in pos correspond to the depth of data. As long as you're dealing with 3D matrices, getting i, j, and k well never get more complicated than this.

On python-3.6+, you can shorten this to—

>>> data[[*pos.T]]
array([6, 2, 4])
cs95
  • 379,657
  • 97
  • 704
  • 746
  • Ouch, downvoter please let me know if there's anything semantically wrong with this answer (from what I can tell, this gives OP what they're expecting). I'd like to correct it! Thanks. – cs95 May 16 '18 at 04:04
  • In a sense, you are doing `[i for i in pos.T]`, except that your solution works only for 3 items. My concern is that the list comprehension seems unnecessary here. Isn't there a native numpy way to reorganize the array? – DYZ May 16 '18 at 04:11
  • @DyZ I added an explanation. If you want to index more items, you will have more rows, not columns. The columns correspond to the depth; and as long as you're dealing with 3D arrays, getting i, j, and k will be done the same way. This will work for 3 items, and it will work for many more. – cs95 May 16 '18 at 04:13
  • 1
    tensorflow has tf.gather_nd. Is there not an analog to numpy? – MFisherKDX May 16 '18 at 04:15
  • @MFisherKDX Tensorflow is built on top of numpy. It likely does something like this under the hood. But I'll keep looking. This is the most straightforward solution I am aware of :) – cs95 May 16 '18 at 04:18
  • @MFisherKDX In fact, gather is what I want. I thought `np.take` would work, but it does not seem to support nd arrays. – DYZ May 16 '18 at 04:18
  • 1
    Yeah. I downvoted this answer and started to post my own using take and then realized it didn't work and removed my downvote. Sorry Coldspeed – MFisherKDX May 16 '18 at 04:20
  • 1
    @MFisherKDX Well, if OP doesn't mind a python-3.6 solution, I've added it. – cs95 May 16 '18 at 04:34
  • `data[[*pos.T]]` is deprecated since [v1.15.0](https://github.com/numpy/numpy/releases/tag/v1.15.0): "Multidimensional indexing with anything but a tuple is deprecated." Use `data[tuple(pos.T)]` instead. – John May 06 '20 at 10:03