來自索引列表的 numpy 花式索引與切片相結合

Question

我有一個大小為 (d1,d2,d3,d4) 的 4 維數組，以及一個包含k個條目的索引列表。 索引列表的每個元素都有 3 個元素。 第一個是數組第一維的索引。 第二個和第三個對應於第 3 個和第 4 個維度的切片的開始，其中切片長度是固定數字n 。 從數組的第二維我總是想要所有元素。 最后我想要一個大小為 (k,d3,n,n) 的新數組。 出於性能原因，我想在沒有 for 循環的情況下執行此操作，並且如果可能，沒有數據復制。 以下解決方案有效，但使用 for 循環並復制數據：

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來獲得滑動窗口。 有關使用基於view_as_windows的as_strided的更多信息。 然后，使用來自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]]