Функция, аналогичная numpy.take, индексирует выходной параметр

Question

В документации numpy.takeзаявлено, что a индексируется в соответствии с indices а также axisзатем результат может быть сохранен в out параметр. Существует ли функция, которая выполняет индексацию out вместо? Используя необычное индексирование, это будет что-то вроде:

out[:, :, indices, :] = a

Здесь я предполагаю, что axis=2 но в моем случае я не знаю оси заранее. Также приемлемо решение с использованием 1d логических масок вместо индексов.

Answer 1

Вы можете использовать swapaxes следующим образом:

>>> A = np.arange(24).reshape(2,3,4)
>>> out = np.empty_like(A)
>>> I = [2,0,1]
>>> axis = 1
>>> out.swapaxes(0, axis)[I] = A.swapaxes(0, axis)
>>> out                                                                                                             
array([[[ 4,  5,  6,  7],                                                                                           
        [ 8,  9, 10, 11],                                                                                           
        [ 0,  1,  2,  3]],                                                                                          

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

Answer 2

Некоторые из numpy Функции создают индексный кортеж при работе на указанной оси.

Код не очень красивый, но общий и достаточно эффективный.

In [700]: out = np.zeros((2,1,4,5),int)
In [701]: out
Out[701]: 
array([[[[0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0]]],


       [[[0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0]]]])

In [702]: indices = [3,0,1]

Сделайте кортеж индексации. Начните со списка или массива для простоты построения, а затем преобразуйте в tuple при индексации:

In [703]: idx = [slice(None)]*out.ndim
In [704]: idx[2] = indices
In [705]: idx
Out[705]: 
[slice(None, None, None),
 slice(None, None, None),
 [3, 0, 1],
 slice(None, None, None)]
In [706]: out[tuple(idx)] = 10    
In [707]: out
Out[707]: 
array([[[[10, 10, 10, 10, 10],
         [10, 10, 10, 10, 10],
         [ 0,  0,  0,  0,  0],
         [10, 10, 10, 10, 10]]],


       [[[10, 10, 10, 10, 10],
         [10, 10, 10, 10, 10],
         [ 0,  0,  0,  0,  0],
         [10, 10, 10, 10, 10]]]])

Это соответствует take:

In [708]: np.take(out, indices, axis=2)
Out[708]: 
array([[[[10, 10, 10, 10, 10],
         [10, 10, 10, 10, 10],
         [10, 10, 10, 10, 10]]],


       [[[10, 10, 10, 10, 10],
         [10, 10, 10, 10, 10],
         [10, 10, 10, 10, 10]]]])

Мы можем установить более сложные значения при условии правильной трансляции:

out[tuple(idx)] = np.array([10,11,12])[...,None]

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

---

Есть такие функции, как place, put, copyto которые предоставляют другие способы управления присваиванием (помимо обычной индексации). Но никто не берет axis параметр как np.take,