使用numpy的`as_strided`函数创建任意维度的补丁、平铺、滚动或滑动窗口

14
今天早上花了一些时间寻找一个通用问题,以便将有关as_strided和/或如何制作广义窗口函数的重复问题指向。看起来有很多关于如何(安全地)创建机器学习、卷积、图像处理和/或数值积分的数组补丁、滑动窗口、滚动窗口、瓦片或视图的问题
我正在寻找一个通用函数,可以接受窗口、步幅和轴参数,并返回任意维度的as_strided视图。我将在下面给出我的答案,但我感兴趣的是是否有人可以制定更有效的方法,因为我不确定使用np.squeeze()是否是最佳方法,我不确定我的assert语句是否足以使该函数安全地写入结果视图,也不知道如何处理轴不按升序排列的边缘情况。
尽职调查:
我能找到的最通用的函数是由@eickenberg编写的sklearn.feature_extraction.image.extract_patches(以及明显等效的skimage.util.view_as_windows),但这些在互联网上并没有得到很好的记录,并且无法对少于原始数组中存在的轴执行窗口操作(例如,这个问题要求在仅一个轴上的一个特定大小的窗口)。还经常会有问题需要一个只用numpy的答案。
@Divakar为1D输入创建了一个通用的numpy函数这里,但高维输入需要更多的注意。我已经制作了一个基本的2D窗口覆盖3D输入方法,但它不太可扩展。

如果您能想到其他人可能搜索的流行词,请随意编辑第一段内容。 - Daniel F
根据我的经验,你应该真正使用skimage.util.view_as_windows()。自己编写并不值得。 - Nils Werner
@NilsWerner 对于二维图像?当然可以。但对于更高维度的数据,你想要在任意轴上使用窗口吗?不太行。我从事机械和机电学的机器学习,我的状态空间可能有超过10个维度,而我可能只想在一个小选择范围内使用窗口函数。 - Daniel F
为什么不呢?它适用于任意维度... - Nils Werner
啊,看起来skimage.util.view_as_windows()已经被升级到nD版本了,自从我写这个东西之后。无论如何,对于商业API来说,numpyskimage更常见(当我最初编写时,因需处理非可扩展的abaqus API而陷入困境),因此它仍然对某些人有用。 - Daniel F
是的,你的解决方案并不无用(我也一直在创建和提出自己的解决方案)。但我也学到了,在大多数情况下,我可以避免所有这些痛苦,只需回退到 skimage 即可。 - Nils Werner
1个回答

15

2020年1月修订:将可迭代返回值从列表更改为生成器以节省内存。

2020年10月修订:将生成器放在单独的函数中,因为混合使用生成器和return语句不直观。

以下是我目前拥有的配方:

def window_nd(a, window, steps = None, axis = None, gen_data = False):
        """
        Create a windowed view over `n`-dimensional input that uses an 
        `m`-dimensional window, with `m <= n`
        
        Parameters
        -------------
        a : Array-like
            The array to create the view on
            
        window : tuple or int
            If int, the size of the window in `axis`, or in all dimensions if 
            `axis == None`
            
            If tuple, the shape of the desired window.  `window.size` must be:
                equal to `len(axis)` if `axis != None`, else 
                equal to `len(a.shape)`, or 
                1
                
        steps : tuple, int or None
            The offset between consecutive windows in desired dimension
            If None, offset is one in all dimensions
            If int, the offset for all windows over `axis`
            If tuple, the steps along each `axis`.  
                `len(steps)` must me equal to `len(axis)`
    
        axis : tuple, int or None
            The axes over which to apply the window
            If None, apply over all dimensions
            if tuple or int, the dimensions over which to apply the window

        gen_data : boolean
            returns data needed for a generator
    
        Returns
        -------
        
        a_view : ndarray
            A windowed view on the input array `a`, or `a, wshp`, where `whsp` is the window shape needed for creating the generator
            
        """
        ashp = np.array(a.shape)
        
        if axis != None:
            axs = np.array(axis, ndmin = 1)
            assert np.all(np.in1d(axs, np.arange(ashp.size))), "Axes out of range"
        else:
            axs = np.arange(ashp.size)
            
        window = np.array(window, ndmin = 1)
        assert (window.size == axs.size) | (window.size == 1), "Window dims and axes don't match"
        wshp = ashp.copy()
        wshp[axs] = window
        assert np.all(wshp <= ashp), "Window is bigger than input array in axes"
        
        stp = np.ones_like(ashp)
        if steps:
            steps = np.array(steps, ndmin = 1)
            assert np.all(steps > 0), "Only positive steps allowed"
            assert (steps.size == axs.size) | (steps.size == 1), "Steps and axes don't match"
            stp[axs] = steps
    
        astr = np.array(a.strides)
        
        shape = tuple((ashp - wshp) // stp + 1) + tuple(wshp)
        strides = tuple(astr * stp) + tuple(astr)
        
        as_strided = np.lib.stride_tricks.as_strided
        a_view = np.squeeze(as_strided(a, 
                                     shape = shape, 
                                     strides = strides))
        if gen_data :
            return a_view, shape[:-wshp.size]
        else:
            return a_view

def window_gen(a, window, **kwargs):
    #Same docstring as above, returns a generator
    _ = kwargs.pop(gen_data, False)
    a_view, shp = window_nd(a, window, gen_data  = True, **kwargs)
    for idx in np.ndindex(shp):
        yield a_view[idx]

一些测试用例:
a = np.arange(1000).reshape(10,10,10)

window_nd(a, 4).shape # sliding (4x4x4) window
Out: (7, 7, 7, 4, 4, 4)

window_nd(a, 2, 2).shape # (2x2x2) blocks
Out: (5, 5, 5, 2, 2, 2)

window_nd(a, 2, 1, 0).shape # sliding window of width 2 over axis 0
Out: (9, 2, 10, 10)

window_nd(a, 2, 2, (0,1)).shape # tiled (2x2) windows over first and second axes
Out: (5, 5, 2, 2, 10)

window_nd(a,(4,3,2)).shape  # arbitrary sliding window
Out: (7, 8, 9, 4, 3, 2)

window_nd(a,(4,3,2),(1,5,2),(0,2,1)).shape #arbitrary windows, steps and axis
Out: (7, 5, 2, 4, 2, 3) # note shape[-3:] != window as axes are out of order

碰巧我正在处理一个非常相似的东西(如果我理解它应该接受的轴/窗口是什么的话,可能是完全一样的)。有一个问题 - 你能展示一个在m < n情况下的形状的样本运行吗? - Divakar
添加了一些测试用例。 - Daniel F
这很有道理!谢谢。之前我尝试过:window_nd(a, window=(3,2), axis=(0,1)),但是出现了错误。而且这正是我一直在处理的任务 :) - Divakar
1
我认为在Python 3.7.7(以及可能更早的版本)中,“generator”关键字存在问题。似乎只要代码中有一个“yield”,它就会返回生成器,即使if语句没有被触发(例如,generator=False也会返回一个迭代器)。 - skjerns
@skjerns,抱歉我花了太长时间来修复。 - Daniel F

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接