2d массив как индекс 3d массива

Question

У меня был двумерный массив (C) с элементами 8000x64, одномерный массив (ы) с элементами 8000x1 и еще один одномерный массив (d) с элементами 1x64. Каждая строка индекса i, где s [i] - True, должна быть добавлена ​​вектором d. Это работает довольно хорошо:

C[s == True] += d

Теперь я добавил одно измерение к C, s и d, и логика c выше должна быть применена к каждому элементу дополнительного измерения.

Следующий код делает то, что я хочу, но он очень медленный.

for i in range(I):
        C_this = C[:,:,i]
        s_this = s[:,i]
        d_this = d[:,i]

        C_this[s_this == True] += d_this
        C[:,:,i] = C_this

Есть ли способ numpy сделать это без for для l oop?

Answer 1

Проще с дополнительным измерением в начале:

In [376]: C = np.zeros((4,2,3),int)                                                            
In [377]: s = np.array([[0,0],[0,1],[1,0],[1,1]],bool)                                         
In [378]: d = np.arange(1,13).reshape(4,3)                                                     
In [379]: C.shape, s.shape, d.shape                                                            
Out[379]: ((4, 2, 3), (4, 2), (4, 3))
In [380]: I,J = np.nonzero(s)                                                                  
In [381]: I,J                                                                                  
Out[381]: (array([1, 2, 3, 3]), array([1, 0, 0, 1]))

In [383]: C[I,J]=d[I]                                                                          
In [384]: C                                                                                    
Out[384]: 
array([[[ 0,  0,  0],
        [ 0,  0,  0]],

       [[ 0,  0,  0],
        [ 4,  5,  6]],

       [[ 7,  8,  9],
        [ 0,  0,  0]],

       [[10, 11, 12],
        [10, 11, 12]]])

Ваш путь:

In [385]: C = np.zeros((4,2,3),int)                                                            
In [386]: for i in range(4): 
     ...:     C[i,:,:][s[i,:]] += d[i,:] 
     ...:                                                                                      
In [387]: C                                                                                    
Out[387]: 
array([[[ 0,  0,  0],
        [ 0,  0,  0]],

       [[ 0,  0,  0],
        [ 4,  5,  6]],

       [[ 7,  8,  9],
        [ 0,  0,  0]],

       [[10, 11, 12],
        [10, 11, 12]]])

Answer 2

В связи с тем, как работает индексирование numpy, s выбирает соответствующие строки C в первом примере. Чтобы сделать то же самое в трехмерном случае, вам придется преобразовать C во что-то, что (8000*3, 64) и s в (8000*3, 1). Теперь единственной проблемой является получение d для учета различного количества строк в каждом третьем измерении, что можно сделать с помощью np.repeat.

Первая часть -

C2 = np.swapaxes(C, -1, 1).reshape(-1, 64)

Это крайне неэффективно, потому что копирует весь ваш массив. Лучшее расположение было бы, если бы C имел форму (3, 8000, 64) для начала. Тогда вам нужно будет только развернуть первые две оси, чтобы получить правильную форму и расположение в памяти, без копирования данных.

repeats = np.count_nonzero(s, axis=0)
C.reshape(-1, 64)[s.ravel()] += np.repeat(d, repeats, axis=0)

Поскольку операция преобразования в этом случае возвращает представление, индексирование должно работать правильно, чтобы приращение на месте. Я не думаю, что этот подход обязательно очень хорош, поскольку он копирует каждую строку d столько раз, сколько s не равно нулю в соответствующем элементе нового измерения.

Answer 3

Вот моя реализация предложенного метода @hpaulj. Обратите внимание, что я не хочу брать у него кредит, поэтому, пожалуйста, пометьте его ответ, а не мой, как правильный. Просто хотел поделиться тем, что я сделал.

import numpy as np
import numpy.random as npr

C = np.zeros((100, 8000, 64), dtype=int)
s = np.zeros((100, 8000), dtype=bool)
d = np.zeros((100, 64), dtype=int)

C[:,:,:] = npr.randint(50, size=C.shape)
s[:,:] = npr.randint(3, size=s.shape)
d[:,:] = npr.randint(10, size=d.shape)

I, J = np.nonzero(s)
C[I, J] += d[I]

Затем я профилировал созданную мной программу, и она запускалась на моей машине менее чем за 450 миллисекунд (последние две строки занимают менее 300 мс). Обратите внимание, что вызовы «randint» были просто для установки значений массива, поэтому эти строки не будут применяться в вашем случае использования.