NumPy: индексирование массива по списку кортежей - как это сделать правильно?

Question

Я нахожусь в следующей ситуации - у меня есть следующее:

• Многомерный массив NumPy a из n Размеры
• t, массив k строки (кортежи), каждая с n элементами. Другими словами, каждая строка в этом массиве является индексом в a

Что я хочу: из a, вернуть массив b с k скалярными элементами, i th-элемент в b является результатом индексации a с i -ым кортежем из t.

Кажется достаточно тривиальным. Следующий подход, однако, не работает

def get(a, t):
    # wrong result + takes way too long
    return a[t]

Я должен прибегнуть к этому итеративно, то есть следующее работает правильно:

def get(a, t):
    res = []
    for ind in t:
        a_scalar = a
        for i in ind:
            a_scalar = a_scalar[i]

        # a_scalar is now a scalar
        res.append(a_scalar)

    return res

Это работает, за исключением того факта, чточто в каждом измерении в a содержится более 30 элементов, процедура становится действительно медленной, когда n достигает более 5. Я понимаю, что это будет медленно, независимо от того, однако я хотел бы использовать возможности numpy, как мне кажетсязначительно ускорит этот процесс.

Answer 1

Ключом к правильному пониманию является понимание роли индексирования списков и кортежей. Часто они обрабатываются одинаково, но в numpy индексирование, кортежи, списки и массивы передают разную информацию.

In [1]: a = np.arange(12).reshape(3,4)                                          
In [2]: t = np.array([(0,0),(1,1),(2,2)])                                       

In [4]: a                                                                       
Out[4]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
In [5]: t                                                                       
Out[5]: 
array([[0, 0],
       [1, 1],
       [2, 2]])

Вы пробовали:

In [6]: a[t]                                                                    
Out[6]: 
array([[[ 0,  1,  2,  3],
        [ 0,  1,  2,  3]],

       [[ 4,  5,  6,  7],
        [ 4,  5,  6,  7]],

       [[ 8,  9, 10, 11],
        [ 8,  9, 10, 11]]])

Так что же с этим не так? ? Он запустился, но выбрал (3,2) массив строк a. То есть он применил t только к первому измерению, фактически a[t, :]. Вы хотите индексировать по всем измерениям, что-то вроде a[t1, t2]. Это то же самое, что и a[(t1,t2)] - кортеж индексов.

In [10]: a[tuple(t[0])]                # a[(0,0)]                                         
Out[10]: 0
In [11]: a[tuple(t[1])]                # a[(1,1)]                                         
Out[11]: 5
In [12]: a[tuple(t[2])]                                                         
Out[12]: 10

или выполнение всего сразу:

In [13]: a[(t[:,0], t[:,1])]                                                      
Out[13]: array([ 0,  5, 10])

Другой способ написать это - n списки (илимассивы), по одному для каждого измерения:

In [14]: a[[0,1,2],[0,1,2]]                                                     
Out[14]: array([ 0,  5, 10])

In [18]: tuple(t.T)                                                             
Out[18]: (array([0, 1, 2]), array([0, 1, 2]))
In [19]: a[tuple(t.T)]                                                          
Out[19]: array([ 0,  5, 10])

В общем случае в a[idx1, idx2] массив idx1 передается против idx2 для создания полного массива выбора. Здесь 2 массива 1d и совпадают, выбор - ваш набор пар t. Но тот же принцип применим к выбору набора строк и столбцов: a[ [[0],[2]], [0,2,3] ].

Используя идеи из [10] и далее, ваш get может быть ускорен:

In [20]: def get(a, t): 
    ...:     res = [] 
    ...:     for ind in t: 
    ...:         res.append(a[tuple(ind)])  # index all dimensions at once 
    ...:     return res 
    ...:                                                                        
In [21]: get(a,t)                                                               
Out[21]: [0, 5, 10]

Если бы t действительно был списком кортежей (в отличие от массива, построенного из них), ваш get мог бы быть:

In [23]: tl = [(0,0),(1,1),(2,2)]                                               
In [24]: [a[ind] for ind in tl]                                                 
Out[24]: [0, 5, 10]

Answer 2

Исследуйте с помощью np.ravel_multi_index

Создайте несколько тестовых данных

arr = np.arange(10**4)
arr.shape=10,10,10,10
t = []
for j in range(5):
    t.append( tuple(np.random.randint(10, size = 4)))

print(t)
# [(1, 8, 2, 0),
#  (2, 3, 3, 6),
#  (1, 4, 8, 5),
#  (2, 2, 6, 3),
#  (0, 5, 0, 2),]

ta = np.array(t).T
print(ta)
# array([[1, 2, 1, 2, 0],
#        [8, 3, 4, 2, 5],
#        [2, 3, 8, 6, 0],
#        [0, 6, 5, 3, 2]])

arr.ravel()[np.ravel_multi_index(tuple(ta), (10,10,10,10))]
# array([1820, 2336, 1485, 2263,  502]

np.ravel_multi_index в основном вычисляет из кортежа входных массивов индекс в уплощенный массив, который начинается с фигуры(в данном случае) (10, 10, 10, 10).

Делает ли это то, что вам нужно? Это достаточно быстро?