Numpy операция по расширению массива в последовательные срезы заданной длины?

Question

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 l oop, но мне было интересно, существует ли метод векторизации 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

Answer 1

Если вы осторожны с математикой и учитываете предупреждения в документах, вы можете использовать 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]])

Answer 2

Как насчет этой функции?

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)