用元组数组进行Numpy 2d索引，可能还进行Nd索引

Question

With 1d arrays it's possible to index it by Nd array of integers like this:

>>> rand = np.random.rand(9).astype(np.float32)
>>> rand
array([ 0.69786191,  0.09376735,  0.60141236,  0.35305005,  0.68340319,
        0.0746202 ,  0.11620298,  0.46607161,  0.90864712], dtype=float32)
>>> u = np.random.randint(0, 9, (2,2))
>>> u
array([[0, 6],
       [5, 6]])
>>> rand[u]
array([[ 0.69786191,  0.11620298],
       [ 0.0746202 ,  0.11620298]], dtype=float32)

But I can't do the same with 2d arrays:

>>> rand2d = np.random.rand(9).astype(np.float32).reshape(3,3)
>>> rand2d
array([[ 0.83248657,  0.75025952,  0.87252802],
       [ 0.78049046,  0.92902303,  0.42035589],
       [ 0.80461669,  0.49386421,  0.56518084]], dtype=float32)
>>> u = np.random.randint(0, 3, (2,2,2))
>>> u
array([[[2, 2],
        [2, 2]],

       [[0, 2],
        [0, 1]]])
>>> rand2d[u]
array([[[[ 0.80461669,  0.49386421,  0.56518084],
         [ 0.80461669,  0.49386421,  0.56518084]],

        [[ 0.80461669,  0.49386421,  0.56518084],
         [ 0.80461669,  0.49386421,  0.56518084]]],

       [[[ 0.83248657,  0.75025952,  0.87252802],
         [ 0.80461669,  0.49386421,  0.56518084]],

        [[ 0.83248657,  0.75025952,  0.87252802],
         [ 0.78049046,  0.92902303,  0.42035589]]]], dtype=float32)

While the result I expected is:

[[rand2d[2, 2], rand2d[2, 2]],
[rand2d[0, 2], rand2d[0, 1]]] ==
[[0.56518084, 0.56518084],
[0.87252802, 0.75025952]]

How can I achieve this without iterating?

Answer 1

Straight from the example in the docs :

>>> 
>>> x
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11]])
>>> 
>>> rows = np.array([[0,0],[3,3]])
>>> columns = np.array([[0,2],[0,2]])
>>> x[rows,columns]
array([[ 0,  2],
       [ 9, 11]])
>>> 

You can see that it is selecting items at (0,0), (0,2) and (3,0),(3,2).

Answer 2

You can have a look at unravel_index . Not sure if this is exactly what you are after, but it might be useful:

import numpy as np

rand2d = np.array([[ 0.83248657,  0.75025952,  0.87252802],
       [ 0.78049046,  0.92902303,  0.42035589],
       [ 0.80461669,  0.49386421,  0.56518084]], dtype=np.float32)


u = np.random.randint(0, 3, (2,2,2))

print(rand2d[np.unravel_index(u, rand2d.shape)])

Example output is:

[[[ 0.87252802  0.75025952]
  [ 0.83248657  0.75025952]]

 [[ 0.87252802  0.87252802]
  [ 0.75025952  0.87252802]]]

Answer 3

Thanx wwii. My point is indexing an array with matrix of each element coordinates (as example - for per-pixel shift). For my case solution is:

>>> u
array([[[2, 2],
        [2, 2]],

       [[0, 2],
        [0, 1]]])
>>> ux = u.transpose(2,0,1)
>>> ux
array([[[2, 2],
        [0, 0]],

       [[2, 2],
        [2, 1]]])
>>> rand2d[ux[0], ux[1]]
array([[ 0.56518084,  0.56518084],
       [ 0.87252802,  0.75025952]], dtype=float32)

Also, here is my solution to get this array of coordinates, operate with it and use it back for indexing:

>>> ux = np.indices(rand2d.shape)
>>> ux
array([[[0, 0, 0],
        [1, 1, 1],
        [2, 2, 2]],

       [[0, 1, 2],
        [0, 1, 2],
        [0, 1, 2]]])
>>> u = ux.transpose(1,2,0)
>>> u
array([[[0, 0],
        [0, 1],
        [0, 2]],

       [[1, 0],
        [1, 1],
        [1, 2]],

       [[2, 0],
        [2, 1],
        [2, 2]]])
>>> u[1,1]-=1
>>> u
array([[[0, 0],
        [0, 1],
        [0, 2]],

       [[1, 0],
        [0, 0],
        [1, 2]],

       [[2, 0],
        [2, 1],
        [2, 2]]])
>>> ux = u.transpose(2,0,1) #Transpose back
>>> ux
array([[[0, 0, 0],
        [1, 0, 1],
        [2, 2, 2]],

       [[0, 1, 2],
        [0, 0, 2],
        [0, 1, 2]]])
>>> rand2d[ux[0], ux[1]]
array([[ 0.83248657,  0.75025952,  0.87252802],
       [ 0.78049046,  0.83248657,  0.42035589],
       [ 0.80461669,  0.49386421,  0.56518084]], dtype=float32)