使用numpy的as_strided檢索中心的主對角線

Question

我有一個正方形陣列的子陣列x，形狀(N, N)，我想檢索這是在主對角線的x居中形狀的方形子陣列。例如，N = 3 & n = 2，並

操作

x = np.arange(9).reshape((3, 3)) 

應該產生

array([[[0, 1], 
     [3, 4]], 
     [[4, 5], 
     [7, 8]]]) 

一種方法是使用make_windows

def make_windows(a, sub_w, sub_h): 
    w, h = a.shape 
    a_strided = np.lib.stride_tricks.as_strided(
     a, shape=[w - sub_w + 1, h - sub_h + 1, 
        sub_w, sub_h], 
    strides=a.strides + a.strides) 
    return a_strided 

，並完成類似np.einsum('ii...->i...', make_windows(x, 2, 2))，但它會整齊地做到一步。單獨使用as_strided可以嗎？

Answer 1

不確定：

def diag_windows(x, n): 
    if x.ndim != 2 or x.shape[0] != x.shape[1] or x.shape[0] < n: 
     raise ValueError("Invalid input") 
    w = as_strided(x, shape=(x.shape[0] - n + 1, n, n), 
        strides=(x.strides[0]+x.strides[1], x.strides[0], x.strides[1])) 
    return w 

例如：

In [14]: x 
Out[14]: 
array([[ 0, 1, 2, 3], 
     [ 4, 5, 6, 7], 
     [ 8, 9, 10, 11], 
     [12, 13, 14, 15]]) 

In [15]: diag_windows(x, 2) 
Out[15]: 
array([[[ 0, 1], 
     [ 4, 5]], 

     [[ 5, 6], 
     [ 9, 10]], 

     [[10, 11], 
     [14, 15]]]) 

In [16]: diag_windows(x, 3) 
Out[16]: 
array([[[ 0, 1, 2], 
     [ 4, 5, 6], 
     [ 8, 9, 10]], 

     [[ 5, 6, 7], 
     [ 9, 10, 11], 
     [13, 14, 15]]]) 

In [17]: diag_windows(x, 4) 
Out[17]: 
array([[[ 0, 1, 2, 3], 
     [ 4, 5, 6, 7], 
     [ 8, 9, 10, 11], 
     [12, 13, 14, 15]]])