有一个绝妙的技巧stride_tricks
,您可以在 SO 和其他上找到具有不同通用性的滚动窗口函数(numpy 本身目前没有),这是根据您所获得的内容量身定制的版本:
def rolling_window(arr, window):
"""Very basic multi dimensional rolling window. window should be the shape of
of the desired subarrays. Window is either a scalar or a tuple of same size
as `arr.shape`.
"""
shape = np.array(arr.shape*2)
strides = np.array(arr.strides*2)
window = np.asarray(window)
shape[arr.ndim:] = window # new dimensions size
shape[:arr.ndim] -= window - 1
if np.any(shape < 1):
raise ValueError('window size is too large')
return np.lib.stride_tricks.as_strided(arr, shape=shape, strides=strides)
# Now:
view = rolling_window(arr, 2)
view[0,0] # first slice in your loop
请注意,它view
包含与原始数组相同的数据!这可能会导致意想不到的结果。但是您似乎只想要对角线,您也可以使用跨步技巧来确保不复制数据(下一个版本将创建一个带有 的视图diagonal
,旧版本总是一个副本):
diagonal = np.diagonal(view, axis1=0, axis2=1)
# unfortunatly now the first slice is diagonal[...,0], so just roll it to the start:
diagonal = np.rollaxis(diagonal, -1)
现在diagonal
是您在 for 循环中创建的数组(在较新的版本中,.copy()
如果您不想要视图,请添加 a)。
编辑:由于slices
数组是 2D 而不是 3D 因为您追加,所以这里缺少重塑:
slices = diagonal.reshape(-1,2)
如果您有这么小的数组,这可能不会更快,但它的常量(期望复制diagonal
调用中的数据)与数组大小有关。