numpy 执行滑动窗口的矢量化方式是什么

cwtwac6a  于 2023-03-23  发布在  其他
关注(0)|答案(1)|浏览(176)

我有一个嵌套的for循环函数。对于2D矩阵的每个索引i和j,它对2D数组的2D切片的所有元素求和,如sum(data[i-1:i+1,j-1+i+1]))。

import numpy as np

data=np.array([[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]])

# This is to specify at the edge indices that the sum wraps around
pad_factor=1
data_padded = np.pad(data, pad_factor, mode='wrap')
print(data_padded)
output:
[[16 13 14 15 16 13]
 [ 4  1  2  3  4  1]
 [ 8  5  6  7  8  5]
 [12  9 10 11 12  9]
 [16 13 14 15 16 13]
 [ 4  1  2  3  4  1]]

result=np.zeros((np.shape(data)))
for i in range(0,np.shape(data)[0]):
    for j in range(0,np.shape(data)[1]):
        result[i,j] =  np.sum(data_padded[i-1+pad_factor:i+1+pad_factor+1, j-1+pad_factor:j+1+pad_factor+1])

print(result)

output:
[[69. 66. 75. 72.]
 [57. 54. 63. 60.]
 [93. 90. 99. 96.]
 [81. 78. 87. 84.]]

然而,在一个较大的数组上,这花费的时间太长了。所以我想把它矢量化。我试着创建一个meshgrid,然后把这些数组输入公式:

i, j = np.mgrid[0:np.shape(data)[0],0:np.shape(data)[1]]
result=np.sum(data_padded[i-1:i+1+1,j-1:j+1+1])

这将产生以下错误:

TypeError: only integer scalar arrays can be converted to a scalar index

它不喜欢将给定数组的数组切片作为输入。
然而,同样的方法适用于矩阵中的单个元素,例如:

i, j = np.mgrid[0:np.shape(data)[0]-1,0:np.shape(data)[1]-1]

result=data[i,j]
print(result)

output
[[ 1  2  3]
 [ 5  6  7]
 [ 9 10 11]]

所以我想知道是否有办法做到这一点。
我也对向量化原始问题的解决方案感兴趣。

ca1c2owp

ca1c2owp1#

这是一个滑动窗口任务。stride_tricks子模块有一些工具可以使用strides来创建多维view。在本例中,我们创建一个(4,4,3,3)视图,并对最后2个维度求和:

In [72]: np.lib.stride_tricks.sliding_window_view(data_padded,(3,3)).sum(axis=(2,3))
Out[72]: 
array([[69, 66, 75, 72],
       [57, 54, 63, 60],
       [93, 90, 99, 96],
       [81, 78, 87, 84]])

编辑

为了简化示例,让我们尝试1d索引

In [93]: x=np.arange(10,100,10);x
Out[93]: array([10, 20, 30, 40, 50, 60, 70, 80, 90])

迭代地,我们可以得到一组3个元素窗口,其中:

In [94]: [x[i:i+3] for i in range(5)]
Out[94]: 
[array([10, 20, 30]),
 array([20, 30, 40]),
 array([30, 40, 50]),
 array([40, 50, 60]),
 array([50, 60, 70])]

但正如你所发现的,切片不能使用数组作为开始/停止值:

In [96]: i = np.arange(0,5); x[i:i+3]
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[96], line 1
----> 1 i = np.arange(0,5); x[i:i+3]

TypeError: only integer scalar arrays can be converted to a scalar index

我们可以创建一个索引数组(不是切片):

In [97]: idx = np.arange(5)[:,None]+np.arange(3)  # np.linspace also works    
In [98]: idx
Out[98]: 
array([[0, 1, 2],
       [1, 2, 3],
       [2, 3, 4],
       [3, 4, 5],
       [4, 5, 6]])    
In [99]: x[idx]
Out[99]: 
array([[10, 20, 30],
       [20, 30, 40],
       [30, 40, 50],
       [40, 50, 60],
       [50, 60, 70]])

In [100]: np.lib.stride_tricks.sliding_window_view(x,3)
Out[100]: 
array([[10, 20, 30],
       [20, 30, 40],
       [30, 40, 50],
       [40, 50, 60],
       [50, 60, 70],
       [60, 70, 80],
       [70, 80, 90]])

In [101]: _.strides
Out[101]: (4, 4)

strides在两个方向上都是4个字节,或者一个元素。其中,x被整形为普通的(3,3)数组,向下移动3个元素:

In [105]: x.reshape(3,3).strides
Out[105]: (12, 4)

相关问题