numpy: широковещательный массив, катящийся вдоль новой оси с переменным смещением, заданным во втором массиве

Question

Я знаю, что numpy.roll может сдвигать массив вдоль одной или нескольких существующих осей. Как мне создать новую ось в массиве x, вдоль которой я хочу, чтобы представления или ее копии были свернуты массивом shift?

Пример:

x = np.arange(10)
shift = np.array([2, 4])

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

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

Редактировать: Я ищу общее решение (в идеале без зацикливания), которое также можно применить к массивам более высокой размерности. Другой пример:

x = np.arange(20).reshape(2, 10)
shift = np.array([2, 4])

#input
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14, 15, 16, 17, 18, 19]])

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

   [[18, 16],
    [19, 17],
    [10, 18],
    [11, 19],
    [12, 10],
    [13, 11],
    [14, 12],
    [15, 13],
    [16, 14],
    [17, 15]]])

Answer 1

Вот векторизованное решение, использующее broadcasting, которое охватывает общие n-dim случаи массива -

np.take(x,(-shift + np.arange(x.shape[-1])[:,None]),axis=-1)

Образцы прогонов

1) x как 1D -

In [114]: x = np.arange(10)
     ...: shift = np.array([2, 4])

In [115]: np.take(x,(-shift + np.arange(x.shape[-1])[:,None]),axis=-1)
Out[115]: 
array([[8, 6],
       [9, 7],
       [0, 8],
       [1, 9],
       [2, 0],
       [3, 1],
       [4, 2],
       [5, 3],
       [6, 4],
       [7, 5]])

2) x как 2D -

In [116]: x = np.arange(20).reshape(2, 10)
     ...: shift = np.array([2, 4])

In [117]: np.take(x,(-shift + np.arange(x.shape[-1])[:,None]),axis=-1)
Out[117]: 
array([[[ 8,  6],
        [ 9,  7],
        [ 0,  8],
        [ 1,  9],
        [ 2,  0],
        [ 3,  1],
        [ 4,  2],
        [ 5,  3],
        [ 6,  4],
        [ 7,  5]],

       [[18, 16],
        [19, 17],
        [10, 18],
        [11, 19],
        [12, 10],
        [13, 11],
        [14, 12],
        [15, 13],
        [16, 14],
        [17, 15]]])

Answer 2

Я почти не хочу предоставлять эту альтернативу, потому что я думаю, что ответ @ BenT прост и логичен

np.array([np.roll(x,sh) for sh in shift]).T
np.stack([np.roll(x,sh) for sh in shift], axis=1)  # may be easier to generalize

но я могу сделать оригинальный x=np.arange(10) чехол с as_strided:

Выполнить все смены:

In [352]: arr = np.lib.stride_tricks.as_strided(np.hstack((x,x)),shape=(10,10), strides=(8,8))
In [353]: arr
Out[353]: 
array([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
       [1, 2, 3, 4, 5, 6, 7, 8, 9, 0],
       [2, 3, 4, 5, 6, 7, 8, 9, 0, 1],
       [3, 4, 5, 6, 7, 8, 9, 0, 1, 2],
       [4, 5, 6, 7, 8, 9, 0, 1, 2, 3],
       [5, 6, 7, 8, 9, 0, 1, 2, 3, 4],
       [6, 7, 8, 9, 0, 1, 2, 3, 4, 5],
       [7, 8, 9, 0, 1, 2, 3, 4, 5, 6],
       [8, 9, 0, 1, 2, 3, 4, 5, 6, 7],
       [9, 0, 1, 2, 3, 4, 5, 6, 7, 8]])

Затем выберите те, которые вы хотите:

In [358]: arr[::-1][shift-1]
Out[358]: 
array([[8, 9, 0, 1, 2, 3, 4, 5, 6, 7],
       [6, 7, 8, 9, 0, 1, 2, 3, 4, 5]])

Я написал и протестировал версию stack с одной попытки, но мне пришлось попробовать несколько вещей, чтобы получить as_strided.

Я бы также предпочел обобщить понимание списка для более высоких измерений.

---

Для вашего 2d x:

np.stack([np.roll(x,sh, axis=1) for sh in shift],2)