Clever Geek Handbook

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

I have a 3D array and I need to "shrink" it along the last axis to get a 2D array. I need to do it like this. For each index value for the first two dimensions, I know the index value for the 3rd dimension from which the value should be taken.

For example, I know that if i1 == 2

and i2 == 7

, then i3 == 11

. This means that out[2,7] = inp[2,7,11]

. This mapping from the first two dimensions to the third is represented in another two-dimensional array. In other words, I have an array where in position 2,7

I have 11

as value.

So my question is how to combine these two arrays (3D and 2D) to get the output array (2D).

+3


source to share


4 answers


I believe this should do it:



for i in range(n):
    for j in range(m):
        k = index_mapper[i][j]
        value = input_3d[i][j][k]
        out_2d[i][j] = value
0


source 


In [635]: arr = np.arange(24).reshape(2,3,4)
In [636]: idx = np.array([[1,2,3],[0,1,2]])


In [637]: I,J = np.ogrid[:2,:3]
In [638]: arr[I,J,idx]
Out[638]: 
array([[ 1,  6, 11],
       [12, 17, 22]])
In [639]: arr
Out[639]: 
array([[[ 0,  1,  2,  3],   # 1
        [ 4,  5,  6,  7],   # 6
        [ 8,  9, 10, 11]],  # ll

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

I,J

are broadcast together to select a set of values ​​(2,3) corresponding to idx

:

In [640]: I
Out[640]: 
array([[0],
       [1]])
In [641]: J
Out[641]: array([[0, 1, 2]])


This is a 3d generalization of the simpler 2d task - selecting one item from each row:

In [649]: idx
Out[649]: 
array([[1, 2, 3],
       [0, 1, 2]])
In [650]: idx[np.arange(2), [0,1]]
Out[650]: array([1, 1])

In fact, we could convert the 3d problem to the 2nd one:

In [655]: arr.reshape(6,4)[np.arange(6), idx.ravel()]
Out[655]: array([ 1,  6, 11, 12, 17, 22])
0


source 


inp = np.random.random((20, 10, 5)) # simulate some input
i1, i2 = np.indices(inp.shape[:2])
i3 = np.random.randint(0, 5, size=inp.shape) # or implement whatever mapping
                                             # you want between (i1,i2) and i3
out = inp[(i1, i2, i3)]

See https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#integer-array-indexing for details

0


source


Using numpy.einsum

This can be achieved with a combination of array indexing and using numpy.einsum

:

>>> numpy.einsum('ijij->ij', inp[:, :, indices])

inp[:, :, indices]

creates a four-dimensional array where, for each of the first two indices (the first two dimensions), all the indices of the array of indices are applied to the third dimension. Since the array of indices is two-dimensional, this results in 4D. However, you only need the indices of the array of indices that match those of the first two dimensions. This is then achieved with a string ijij->ij

. This says einsum

that you only want to select items where the 1st, 3rd and 2nd and 4th axis indices are similar. Since the last two dimensions (3rd and 4th) were added by an array of indices, it is like choosing only the index index[i, j]

for the third dimension inp

.

Note that this method can really blow up memory consumption. Especially if inp.shape[:2]

much more inp.shape[2]

, then it inp[:, :, indices].size

will be approximately inp.size ** 2

.

Building indexes manually

First, we create a new array of indices:

>>> idx = numpy.array(list(
...     numpy.ndindex(*inp.shape[:2], 1)  # Python 3 syntax
... ))

Then we update the column corresponding to the third axis:

>>> idx[:, 2] = indices[idx[:, 0], idx[:, 1]]

Now we can select the elements and just change the shape of the result:

>>> inp[tuple(idx.T)].reshape(*inp.shape[:2])

Using numpy.choose

Note: numpy.choose

allows a maximum size of 32 for the axis selected from.

According to this answer and documentation, numpy.choose

we can also use the following:

# First we need to bring the last axis to the front because
# `numpy.choose` chooses from the first axis.
>>> new_inp = numpy.moveaxis(inp, -1, 0)
# Now we can select the elements.
>>> numpy.choose(indices, new_inp)

Although the documentation prohibits using one array for the second argument (choice)

To reduce the likelihood of misinterpretation, even though the following "abuse" is nominally supported, the selection should not or should not be treated as a single array, that is, the outermost sequence-like container should be either a list or a tuple.

this seems to be the case only to prevent misunderstandings:

: sequence of arrays

Selection of arrays. a and all variants must be translated into the same form. If the selection itself is an array (not recommended), then its outer size (ie the one that matches choices.shape[0]

) is taken as defining the "sequence".

So, from my point of view, there is nothing wrong with using numpy.choose

in this way, as long as someone knows what they are doing.

0


source






More articles:


All Articles