выберите все размеры ndarray, но один

Question

У меня есть ndarraym с размером (2,3,4). Я хочу, чтобы выбрать конкретный индекс измерения и суммировать все остальное, как это:

m[:, 2, :].sum() 

но мне нужна функция для обработки выбора элементов, так как ndarray может иметь различные размеры и сумма должна быть в различных срезах. Функция примет размер (здесь 0, 1 или 2), а также индекс (0, 1 или 2) и m[:, 2, :].sum(), если вход (1, 2). Если это (2, 0), я бы ожидал, что они имеют эквивалент m[:, :, 0].sum().

Я не знаю, как написать такую ​​функцию, но что-то, что позволило бы мне поместить ломтик «:» в переменную, может быть? До сих пор я не смог найти решение. Благодаря!

Answer 1

Вы можете поместить ломтик : в переменную с sl = slice(None, None, None). Однако NumPy имеет сокращенный для этого, np.s_[:]

В общем то:

def sum_axis_i(arr, axis, i): 
    idx = (np.s_[:],) * axis + (i,) 
    return arr[idx].sum() 

Обратите внимание, что замыкающие : s является избыточным, поэтому мы не потрудились их создания.

Если вы используете np.index_exp, который очень похож на np.s_, вы можете написать, что:

def sum_axis_i(arr, axis, i): 
    idx = np.index_exp[:] * axis + np.index_exp[i] 
    return arr[idx].sum() 

Answer 2

Вы можете сделать это динамически создавая tuple из slice:

def sum_layer_axis(array, layer, axis): 
    slicer = tuple(slice(None) if ax != axis else layer for ax in range(array.ndim)) 
    return np.sum(array[slicer]) 

Просто для иллюстрации как это работает, рассмотрим то, что интерпретируется из slice:

class Fun(object): 
    def __getitem__(self, item): 
     return item 
Fun()[:, 2, :] 
# (slice(None, None, None), 2, slice(None, None, None)) 

Так что нам нужно было бы воссоздать это (но нам не нужно указывать 3 раза None, одного достаточно). Это возможно с помощью выражения генератора:

tuple(slice(None) if ax != axis else layer for ax in range(array.ndim)) 

поэтому, когда ось не разыскиваемая ось просто вставить : эквивалент: slice(None) и в противном случае вставить указанный слой.

Чтобы проверить результат:

a = np.arange(3*4*5).reshape((3,4,5)) 

axis = 1 
layer = 2 
a[:, layer, :].sum() 
# 480 
sum_layer_axis(a, layer, axis) 
# 480 

Answer 3

Это не может быть наиболее эффективным, но он должен работать:

m = np.arange(60).reshape((3,4,5)) 

def my_sum(m, dim, index): 
    dims = tuple(i for i in xrange(m.ndim) if i != dim) 
    return m.sum(axis=dims)[index] 

print my_sum(m,1,2) 
print m[:,2,:].sum() 

Answer 4

Вот такой подход, который меняет местами первую ось с осью входа, а затем индекс с входным индексом, который выберет все элементы в контексте и, наконец, суммирует все те, которые для конечного результата. Реализация будет компактным с np.swapaxes и , как так -

def sum_axis_index(m, axis, index):  
    return m.swapaxes(0,axis)[index].sum()