Индексирование NumPy с переменным положением

Question

У меня есть массив input_data формы (A, B, C) и массив ind формы (B,). Я хочу пройти через ось B и взять сумму элементов C [B [i]] и C [B [i] +1]. Желаемый выход имеет форму (A, B). У меня есть следующий код, который работает, но я чувствую, что он неэффективен из-за циклизации по индексу через ось B. Есть ли более эффективный метод?

import numpy as np 

input_data = np.random.rand(2, 6, 10) 
ind = [ 2, 3, 5, 6, 5, 4 ] 

out = np.zeros((input_data.shape[0], input_data.shape[1])) 

for i in range(len(ind)): 
    d = input_data[:, i, ind[i]:ind[i]+2] 
    out[:, i] = np.sum(d, axis = 1) 

Edited на основе ответа Divakar в:

import timeit 
import numpy as np 

N = 1000 

input_data = np.random.rand(10, N, 5000) 
ind = (4999 * np.random.rand(N)).astype(np.int) 

def test_1(): # Old loop-based method 
    out = np.zeros((input_data.shape[0], input_data.shape[1])) 

    for i in range(len(ind)): 
     d = input_data[:, i, ind[i]:ind[i]+2] 
     out[:, i] = np.sum(d, axis = 1) 
    return out 

def test_2(): 
    extent = 2 # Comes from 2 in "ind[i]:ind[i]+2" 

    m,n,r = input_data.shape 
    idx = (np.arange(n)*r + ind)[:,None] + np.arange(extent) 
    out1 = input_data.reshape(m,-1)[:,idx].reshape(m,n,-1).sum(2) 
    return out1 

print timeit.timeit(stmt = test_1, number = 1000) 
print timeit.timeit(stmt = test_2, number = 1000) 

print np.all(test_1() == test_2(), keepdims = True) 

>> 7.70429363482 
>> 0.392034666757 
>> [[ True]] 

Answer 1

Вот Векторизованный подход с использованием linear indexing с некоторой помощью broadcasting. Мы объединяем последние две оси входного массива, вычисляем линейные индексы, соответствующие двум последним осям, выполняем срез и преобразуем обратно в трехмерную форму. Наконец, мы суммируем вдоль последней оси, чтобы получить желаемый результат. Реализация будет выглядеть примерно так -

extent = 2 # Comes from 2 in "ind[i]:ind[i]+2" 

m,n,r = input_data.shape 
idx = (np.arange(n)*r + ind)[:,None] + np.arange(extent) 
out1 = input_data.reshape(m,-1)[:,idx].reshape(m,n,-1).sum(2) 

---

Если extent всегда будет 2, как указано в вопросе - "... sum of elements C[B[i]] and C[B[i]+1]", то вы можете просто сделать -

m,n,r = input_data.shape 
ind_arr = np.array(ind) 
axis1_r = np.arange(n) 
out2 = input_data[:,axis1_r,ind_arr] + input_data[:,axis1_r,ind_arr+1] 

Answer 2

Вы могли бы также использовать integer array indexing в сочетании с basic slicing:

import numpy as np 

m,n,r = 2, 6, 10 
input_data = np.arange(2*6*10).reshape(m, n, r) 
ind = np.array([ 2, 3, 5, 6, 5, 4 ]) 
out = np.zeros((input_data.shape[0], input_data.shape[1])) 
for i in range(len(ind)): 
    d = input_data[:, i, ind[i]:ind[i]+2] 
    out[:, i] = np.sum(d, axis = 1) 


out2 = input_data[:, np.arange(n)[:,None], np.add.outer(ind,range(2))].sum(axis=-1) 
print(out2) 
# array([[ 5, 27, 51, 73, 91, 109], 
#  [125, 147, 171, 193, 211, 229]]) 

assert np.allclose(out, out2)