Составление массива с перекрывающимися индексами. Ищем векторизованный путь над петлями

Question

Я ищу векторизованный способ l oop над индексами массива для вертикального объединения их в группы с перекрывающимися индексами.
Чтобы дать суть того, чего я пытаюсь достичь:

Учитывая список [1,2,3,4,5,6] , переменная интервала значения 2 и переменная перекрытия значения 1. Выходные данные должно выглядеть примерно так: [[1,2], [2,3], [3,4], [4,5], [5,6]]

Однако данные, которые у меня есть, имеют форму 1560x2x87236 , где 1560 - объекты, 2x87236 - x, y траектории. Так что для каждого предмета у меня 87236 баллов и 87326 баллов. Поддержание измерения 2, которое представляет xs и ys посредством преобразования, имеет решающее значение.

---

Чтобы упростить представление:

Предположим, у меня есть ndarray:

arr

array([[[35, 33, 34, 42, 32, 30],
        [22, 38, 29, 33, 25, 14]],
       [[17, 25, 39, 17, 41, 22],
        [22, 13, 14, 31, 20, 38]],
       [[30, 10, 33, 25, 38, 26],
        [28, 27, 19, 27, 43, 13]]])

arr.shape

(3, 2, 6)

Я пытаюсь собрать этот массив в группы или интервалы по 3 с перекрывающимися индексами (перекрывающимися на 1 индекс). Вывод выглядит примерно так:

stacked_arr

array([[[ 0.,  0.,  0.],
        [ 0.,  0.,  0.]],

       [[35., 33., 34.],
        [22., 38., 29.]],

       [[34., 42., 32.],
        [29., 33., 25.]],

       [[17., 25., 39.],
        [22., 13., 14.]],

       [[39., 17., 41.],
        [14., 31., 20.]],

       [[30., 10., 33.],
        [28., 27., 19.]],

       [[33., 25., 38.],
        [19., 27., 43.]]])

stacked_arr.shape

(7, 2, 3)

Это Функция, которую я написал, достигает вышеуказанного результата:

def overlap_stack(data, padwith, interv, overlapby):
    sub = 0

    # Initialise: 1 bcuz for a sub, 2 bcuz of x,y
    stacked = cp.zeros(shape=(1, 2, interv))
    while sub < data.shape[0]:
        idx: int
        for idx in range(0, data.shape[2], interv - overlapby):

            # grouping with overlaps
            stack = cp.expand_dims(data[sub, :, idx: idx + interv], axis=0)

            # pad to cope with unequal length
            if (stack.shape[2]) < interv:
                stack = cp.pad(stack, ((0, 0), (0, 0), (0, interv - stack.shape[2])), 'constant',
                               constant_values=padwith)

            # stacking all together
            stacked = cp.vstack((stacked, stack))


        sub += 1
    return stacked

Для преобразования массива 1560x2x87236 требуется от 8 до 10 часов. Буду очень признателен, если вы поможете мне ускорить этот процесс.

Answer 1

Не знаю, знакомы ли вы с numpy .lib.stride_tricks.as_strided , но вот решение, использующее его:

import numpy as np
from numpy.lib.stride_tricks import as_strided

def overlap_stack(data, interv, overlapby):
    A = np.vstack(data)

    window_size = (data.shape[1], interv)
    strides = (window_size[0], interv - overlapby)

    output_strides = (strides[0]*A.strides[0], strides[1]*A.strides[1]) + A.strides

    output_shape = ((A.shape[0] - window_size[0])//strides[0] + 1,
                    (A.shape[1] - window_size[1])//strides[1] + 1) + window_size

    return as_strided(A, shape=output_shape, strides=output_strides).reshape(-1, *output_shape[2:])

Я проигнорировал заполнение потому что я не уверен, как вы этого хотите (хотя вы можете добавить его самостоятельно).

Например:

data = np.array([[[35, 33, 34, 42, 32, 30],
                  [22, 38, 29, 33, 25, 14]],
                 [[17, 25, 39, 17, 41, 22],
                  [22, 13, 14, 31, 20, 38]],
                 [[30, 10, 33, 25, 38, 26],
                  [28, 27, 19, 27, 43, 13]]])

overlap_stack(data, 3, 1)

array([[[35, 33, 34],
        [22, 38, 29]],

       [[34, 42, 32],
        [29, 33, 25]],

       [[17, 25, 39],
        [22, 13, 14]],

       [[39, 17, 41],
        [14, 31, 20]],

       [[30, 10, 33],
        [28, 27, 19]],

       [[33, 25, 38],
        [19, 27, 43]]])

Обратите внимание, что для массива формы (1560, 2, 87236) это будет довольно быстро, но занимает много памяти.