Numpy中任意模板切片的简洁实现方法问询

问题描述 #

在Numpy中，是否存在简洁语法用于创建任意数量相邻数组元素的引用？

当模板宽度硬编码时语法相对简单，例如宽度为3的模板实现如下：

import numpy as np

x = np.arange(8)

# Hard-coded stencil width of 3
x_neighbours = (
    x[  :-2],
    x[ 1:-1],
    x[ 2:  ]
)

但我尝试的任意宽度模板实现可读性较差：

nStencil = 3
x_neighbours = (
   x[indexStart:indexStop]
   for indexStart, indexStop in zip(
      (None, *range(1,nStencil)),
      (*range(1-nStencil,0), None),
   )
)

请问是否有更优的实现方式？

解决方案 #

推荐方法：使用sliding_window_view（Numpy 1.20+） #

Numpy 1.20及以上版本提供了np.lib.stride_tricks.sliding_window_view，这是官方专门用于生成滑动窗口视图的工具，语法简洁且可读性极强：

import numpy as np

x = np.arange(8)
nStencil = 3

# 生成滑动窗口后转置，得到与硬编码版本结构一致的结果
x_neighbours = np.lib.stride_tricks.sliding_window_view(x, window_size=nStencil).T

运行后，x_neighbours是一个形状为(3,6)的Numpy数组，每一行对应原数组的一个切片（和硬编码的x[:-2]、x[1:-1]、x[2:]完全对应）。如果需要保持生成器的形式，可以直接遍历转置后的数组：

x_neighbours = (row for row in np.lib.stride_tricks.sliding_window_view(x, nStencil).T)

这个方法的核心优势：

• 语义清晰，直接表达“滑动窗口”的业务意图
• 基于数组 stride 实现，不会复制原数组数据，执行效率极高
• 支持任意窗口宽度，无需手动计算切片的起始/结束索引

兼容旧版本Numpy：使用as_strided #

如果你的Numpy版本低于1.20，可以用np.lib.stride_tricks.as_strided手动构造视图，但需要自己计算数组的stride参数：

import numpy as np

x = np.arange(8)
nStencil = 3

# 获取数组元素的字节步长
elem_stride = x.strides[0]
# 构造滑动窗口视图
x_neighbours = np.lib.stride_tricks.as_strided(
    x,
    shape=(nStencil, len(x) - nStencil + 1),
    strides=(elem_stride, elem_stride)
)

内容的提问来源于stack exchange，提问作者DavidJ