python - 使用 as_strided 分割 numpy 数组

Question

我正在寻找一种将 numpy 数组分割成重叠 block 的有效方法。我知道 numpy.lib.stride_tricks.as_strided 可能是要走的路，但我似乎无法理解它在适用于任意形状数组的通用函数中的用法。 Here are some examples for specific applications of 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 的跨步 View ，并且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的评论和他的 answer here ，我这样实现了我的功能:

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

这适用于我的一维分割案例，但是，我建议任何寻找适用于任意维度分割的方法的人查看链接答案中的代码。