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

时间:2017-12-18 01:58:50

标签: python numpy

我有一个正方形数组x,形状(N, N),我想检索形状(n, n)的方形子数组,这些子数组以{{1}的主对角线为中心}。例如,使用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))是否可行?

1 个答案:

答案 0 :(得分:3)

不确定

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]]])