Array of indexes for each element alongs the first dimension in a 2D array (numpy., tensorflow)

问题

indexes = np.array([[0,1,3],[1,2,4 ]])
data = np.random.rand(2,5)

Now, i would like an array of shape (2,3), where

result[0] = data[0,indexes[0]]
result[1] = data[1,indexes[1]]

What would be the proper way to achieve this? A numpy way that yould generalize to bigger arrays (perhaps even higher dimensional).

Please note the difference to questions like this, where the array of indexes contains tuples. This is not what I am asking.

Edit

A more general formulation of the question would be:

data.shape == (s0, s1, .., sn)
indexes.shape == (s0, s1, ..., sn-1, K)
so, they have all dimension but the last equal

Than

result[i, j, ..., k] = data[i, j,...,k, indexes[i, j, ..., k]]

where

len([i, j, ..., k]) == len(data)-1 == len(indexes) - 1

回答1:

Here are NumPy and TensorFlow solutions:

import numpy as np
import tensorflow as tf

def gather_index_np(data, index):
    data = np.asarray(data)
    index = np.asarray(index)
    # Make open grid of all but last dimension indices
    grid = np.ogrid[tuple(slice(s) for s in index.shape[:-1])]
    # Add extra dimension in grid
    grid = [g[..., np.newaxis] for g in grid]
    # Complete index
    index_full = tuple(grid + [index])
    # Index data to get result
    result = data[index_full]
    return result

def gather_index_tf(data, index):
    data = tf.convert_to_tensor(data)
    index = tf.convert_to_tensor(index)
    index_shape = tf.shape(index)
    d = index.shape.ndims
    # Make grid of all dimension indices
    grid = tf.meshgrid(*(tf.range(index_shape[i]) for i in range(d)), indexing='ij')
    # Complete index
    index_full = tf.stack(grid[:-1] + [index], axis=-1)
    # Index data to get result
    result = tf.gather_nd(data, index_full)
    return result

Example:

import numpy as np
import tensorflow as tf

data = np.arange(10).reshape((2, 5))
index = np.array([[0, 1, 3], [1, 2, 4]])
print(gather_index_np(data, index))
# [[0 1 3]
#  [6 7 9]]
with tf.Session() as sess:
    print(sess.run(gather_index_tf(data, index)))
# [[0 1 3]
#  [6 7 9]]

回答2:

numpy has take_along_axis which does what you describe plus it also lets you choose the axis.

Example:

>>> a = np.arange(24).reshape(2,3,4)
>>> i = np.random.randint(0,4,(2,3,5))
>>> i
array([[[3, 3, 0, 1, 3],
        [3, 1, 0, 3, 3],
        [3, 2, 0, 3, 3]],

       [[2, 3, 0, 0, 0],
        [1, 1, 3, 1, 2],
        [1, 3, 0, 0, 2]]])

>>> np.take_along_axis(a, i, -1)
array([[[ 3,  3,  0,  1,  3],
        [ 7,  5,  4,  7,  7],
        [11, 10,  8, 11, 11]],

       [[14, 15, 12, 12, 12],
        [17, 17, 19, 17, 18],
        [21, 23, 20, 20, 22]]])

来源：https://stackoverflow.com/questions/53652184/array-of-indexes-for-each-element-alongs-the-first-dimension-in-a-2d-array-nump

标签

python

numpy

tensorflow

indexing