Срез numpy ndarry произвольного размера в массив 1d с заданным списком индексов

#python #numpy #slice #numpy-ndarray #numpy-slicing

Вопрос:

У меня есть numpy ndarray arr , а также indices список индексов, указывающих конкретную запись. Для конкретности возьмем:

 arr = np.arange(2*3*4).reshape((2,3,4))
indices= [1,0,3]
 

У меня есть код, позволяющий выполнять 1d-срезы, arr наблюдая за всеми индексами, кроме одного n :

 arr[:, indices[1], indices[2]]  # n = 0
arr[indices[0], :, indices[2]]  # n = 1
arr[indices[0], indices[1], :]  # n = 2
 

Я хотел бы изменить свой код, чтобы он повторялся n и поддерживал arr произвольное измерение.

Я просмотрел запись о процедурах индексирования в документации и нашел информацию об slice() и np.s_() . Я смог взломать что-то, что работает так, как я хочу:

 def make_custom_slice(n, indices):
    s = list()
    for i, idx in enumerate(indices):
        if i == n:
            s.append(slice(None))
        else:
            s.append(slice(idx, idx 1))
    return tuple(s)


for n in range(arr.ndim):
    np.squeeze(arr[make_custom_slice(n, indices)])
 

Где np.squeeze используется для удаления осей длины 1. Без этого массив, созданный этим, имеет форму (arr.shape[n],1,1,...) , а не (arr.shape[n],) .

Есть ли более идиоматичный способ выполнить эту задачу?

Комментарии:

1. «Есть ли более идиоматичный способ выполнить эту задачу?» Идиоматические способы выполнения задачи с помощью Numpy обычно включают в себя не написание циклов .

2. Я согласен с этим. Улучшения, которые я опубликовал ниже, удаляют цикл в функции, которая создает объект среза. Я не уверен, что функцию можно векторизовать в n, чтобы удалить цикл, в котором она вызывается.

Ответ №1:

Некоторые улучшения в приведенном выше решении (все еще может быть однострочное или более эффективное решение):

 def make_custom_slice(n, indices):
    s = indices.copy()
    s[n] = slice(None)
    return tuple(s)


for n in range(arr.ndim):
    print(arr[make_custom_slice(n, indices)])
 

Целое значение idx может быть использовано для замены объекта среза slice(idx, idx 1) . Поскольку большинство индексов копируются напрямую, начните с копии индексов, а не создавайте список с нуля.

При построении таким образом результат arr[make_custom_slice(n, indices) имеет ожидаемый размер и np.squeeze не является необходимым.