使用as_strided分割numpy数组

Question

我正在寻找一种将numpy数组分割成重叠块的有效方法.我知道这numpy.lib.stride_tricks.as_strided可能是要走的路,但我似乎无法绕过它在一个广义函数中的用法,该函数适用于任意形状的数组.以下是一些特定应用的例子as_strided.

这就是我想要的:

import numpy as np
from numpy.lib.stride_tricks import as_strided

def segment(arr, axis, new_len, step=1, new_axis=None):
    """ Segment an array along some axis.

    Parameters
    ----------
    arr : array-like
        The input array.

    axis : int
        The axis along which to segment.

    new_len : int
        The length of each segment.

    step : int, default 1
        The offset between the start of each segment.

    new_axis : int, optional
        The position where the newly created axis is to be inserted. By
        default, the axis will be added at the end of the array.

    Returns
    -------
    arr_seg : array-like
        The segmented array.
    """

    # calculate shape after segmenting
    new_shape = list(arr.shape)
    new_shape[axis] = (new_shape[axis] - new_len + step) // step
    if new_axis is None:
        new_shape.append(new_len)
    else:
        new_shape.insert(new_axis, new_len)

    # TODO: calculate new strides
    strides = magic_command_returning_strides(...)

    # get view with new strides
    arr_seg = as_strided(arr, new_shape, strides)

    return arr_seg.copy()

所以我想指定要切割成段的轴,段的长度以及它们之间的步长.另外,我想将新轴插入的位置作为参数传递.唯一缺少的是计算步幅.

我知道这可能不会直接用这种方式工作as_strided,即我可能需要实现一个子程序,该子程序返回一个带有step=1和new_axis处于固定位置的跨步视图,然后用所需的step和转置切片.

这是一段可行的代码,但显然很慢:

def segment_slow(arr, axis, new_len, step=1, new_axis=None):
    """ Segment an array along some axis. """

    # calculate shape after segmenting
    new_shape = list(arr.shape)
    new_shape[axis] = (new_shape[axis] - new_len + step) // step
    if new_axis is None:
        new_shape.append(new_len)
    else:
        new_shape.insert(new_axis, new_len)

    # check if the new axis is inserted before the axis to be segmented
    if new_axis is not None and new_axis <= axis:
        axis_in_arr_seg = axis + 1
    else:
        axis_in_arr_seg = axis

    # pre-allocate array
    arr_seg = np.zeros(new_shape, dtype=arr.dtype)

    # setup up indices
    idx_old = [slice(None)] * arr.ndim
    idx_new = [slice(None)] * len(new_shape)

    # get order of transposition for assigning slices to the new array
    order = list(range(arr.ndim))
    if new_axis is None:
        order[-1], order[axis] = order[axis], order[-1]
    elif new_axis > axis:
        order[new_axis-1], order[axis] = order[axis], order[new_axis-1]

    # loop over the axis to be segmented
    for n in range(new_shape[axis_in_arr_seg]):
        idx_old[axis] = n * step + np.arange(new_len)
        idx_new[axis_in_arr_seg] = n
        arr_seg[tuple(idx_new)] = np.transpose(arr[idx_old], order)

    return arr_seg

这是对基本功能的测试:

import numpy.testing as npt    

arr = np.array([[0, 1, 2, 3],
                [4, 5, 6, 7],
                [8, 9, 10, 11]])

arr_seg_1 = segment_slow(arr, axis=1, new_len=3, step=1)
arr_target_1 = np.array([[[0, 1, 2], [1, 2, 3]],
                         [[4, 5, 6], [5, 6, 7]],
                         [[8, 9, 10], [9, 10, 11]]])

npt.assert_allclose(arr_target_1, arr_seg_1)

arr_seg_2 = segment_slow(arr, axis=1, new_len=3, step=1, new_axis=1)
arr_target_2 = np.transpose(arr_target_1, (0, 2, 1))

npt.assert_allclose(arr_target_2, arr_seg_2)

arr_seg_3 = segment_slow(arr, axis=0, new_len=2, step=1)
arr_target_3 = np.array([[[0, 4], [1, 5], [2, 6], [3, 7]],
                         [[4, 8], [5, 9], [6, 10], [7, 11]]])

npt.assert_allclose(arr_target_3, arr_seg_3)

任何帮助将不胜感激!

Answer 1

根据DanielF的评论和他的回答，我实现了我的功能，如下所示：

def segment(arr, axis, new_len, step=1, new_axis=None, return_view=False):
    """ Segment an array along some axis.

    Parameters
    ----------
    arr : array-like
        The input array.

    axis : int
        The axis along which to segment.

    new_len : int
        The length of each segment.

    step : int, default 1
        The offset between the start of each segment.

    new_axis : int, optional
        The position where the newly created axis is to be inserted. By
        default, the axis will be added at the end of the array.

    return_view : bool, default False
        If True, return a view of the segmented array instead of a copy.

    Returns
    -------
    arr_seg : array-like
        The segmented array.
    """

    old_shape = np.array(arr.shape)

    assert new_len <= old_shape[axis],  \
        "new_len is bigger than input array in axis"
    seg_shape = old_shape.copy()
    seg_shape[axis] = new_len

    steps = np.ones_like(old_shape)
    if step:
        step = np.array(step, ndmin = 1)
        assert step > 0, "Only positive steps allowed"
        steps[axis] = step

    arr_strides = np.array(arr.strides)

    shape = tuple((old_shape - seg_shape) // steps + 1) + tuple(seg_shape)
    strides = tuple(arr_strides * steps) + tuple(arr_strides)

    arr_seg = np.squeeze(
        as_strided(arr, shape = shape, strides = strides))

    # squeeze will move the segmented axis to the first position
    arr_seg = np.moveaxis(arr_seg, 0, axis)

    # the new axis comes right after
    if new_axis is not None:
        arr_seg = np.moveaxis(arr_seg, axis+1, new_axis)
    else:
        arr_seg = np.moveaxis(arr_seg, axis+1, -1)

    if return_view:
        return arr_seg
    else:
        return arr_seg.copy()

这对于我的一维段的情况很有效，但是，我建议任何人寻找一种适用于任意维段的方法来检查链接答案中的代码。