python - 使用 numpy as_strided 检索以主对角线为中心的子数组

Question

我有一个方形数组x，形状为(N, N)，我想检索形状为(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]]])