numpy Необычное индексирование из списка индексов в сочетании с ломтиками

Question

У меня есть 4-мерный массив с размером (d1, d2, d3, d4) и список индексов с k записями. Каждый элемент списка индексов имеет 3 элемента. Первый - это индекс для 1-го измерения массива. Второй и третий соответствуют началу срезов для 3-го и 4-го измерения, где длина среза является фиксированным числом n. Из второго измерения массива мне всегда нужны все элементы. В конце я хочу получить новый массив размером (k, d3, n, n). Из соображений производительности я хочу сделать это без for l oop и, если возможно, без копирования данных. Следующее решение работает, но использует для l oop и копирует данные:

d1,d2,d3,d4 = 3,4,10,10
data = np.arange(d1*d2*d3*d4).reshape((d1,d2,d3,d4))
idcs_all = [[0,4,6],
        [2,2,5]]
n = 3
k=len(idcs_all)
sub = np.zeros((k,d2,n,n))

for i, idcs in enumerate(idcs_all):
    sub[i] = data[idcs[0],:, idcs[1]:idcs[1]+n, idcs[2]:idcs[2]+n]

Есть ли способ сделать то же самое с причудливым индексированием?

Answer 1

Мы можем использовать np.lib.stride_tricks.as_strided на основе scikit-image's view_as_windows, чтобы скользить windows. Подробнее об использовании as_strided на основе view_as_windows. Затем, индексирование в эти windows с расширенным индексированием с использованием этих индексов из idcs_all решает это для нас!

from skimage.util import view_as_windows

w = view_as_windows(data,(1,1,n,n))[...,0,0,:,:]
idcs_all = np.asarray(idcs_all)
sub = w[idcs_all[:,0],:,idcs_all[:,1],idcs_all[:,2]]