将数组扩展为给定长度的连续切片的 Numpy 操作?

Numpy operation to expand array into sequential slices of given length?

my_function 必须将 1D numpy 数组扩展为 2D numpy 数组,第二个轴包含 length 的切片,从第一个索引开始直到结束。示例:

import numpy as np
a = np.arange(10)
print (my_function(a, length=3))

预期输出

array([[0, 1, 2],
       [1, 2, 3],
       [2, 3, 4],
       [3, 4, 5],
       [4, 5, 6],
       [5, 6, 7],
       [6, 7, 8],
       [7, 8, 9]])

我可以使用 for 循环实现此目的,但我想知道是否有用于此目的的 numpy 向量化技术。

def my_function(a, length):
    b = np.zeros((len(a)-(length-1), length))
    for i in range(len(b)):
        b[i] = a[i:i+length]
    return b

这个功能怎么样?

import numpy as np
def my_function(a, length):
    result = []
    for i in range(length):
        result.append(a + i)
    return np.vstack(result).T[:len(a) - length + 1]

a = np.arange(10)
length = 3
my_function(a, length)

如果您对数学很认真并注意文档中的警告,则可以使用 np.lib.stride_tricks.as_strided()。您需要为数组计算正确的尺寸,以免溢出。另请注意 as_strided() 共享内存,因此您将在最终输出中多次引用同一内存。 (你当然可以,将其复制到一个新数组)。

>> import numpy as np

>> def my_function(a, length):
    stride = a.strides[0]
    l = len(a) - length + 1
    return np.lib.stride_tricks.as_strided(a, (l, length), (stride,stride) )

>> np.array(my_function(np.arange(10), 3))

array([[0, 1, 2],
       [1, 2, 3],
       [2, 3, 4],
       [3, 4, 5],
       [4, 5, 6],
       [5, 6, 7],
       [6, 7, 8],
       [7, 8, 9]])

>> np.array(my_function(np.arange(15), 7))

array([[ 0,  1,  2,  3,  4,  5,  6],
       [ 1,  2,  3,  4,  5,  6,  7],
       [ 2,  3,  4,  5,  6,  7,  8],
       [ 3,  4,  5,  6,  7,  8,  9],
       [ 4,  5,  6,  7,  8,  9, 10],
       [ 5,  6,  7,  8,  9, 10, 11],
       [ 6,  7,  8,  9, 10, 11, 12],
       [ 7,  8,  9, 10, 11, 12, 13],
       [ 8,  9, 10, 11, 12, 13, 14]])