python: сдвигать каждую матрицу внутри тензора независимо

Question

Этот вопрос очень похож на матрицы Прокрутка строк матрицы независимо

Но я не могу адаптировать его для 3D-тензоров

Мне дали тензор

0 0 0
1 1 1
0 0 0

0 0 0
1 1 1
0 0 0 

и вектор, который определяет, на сколько я хочу сместить мой столбец матриц,

1 2

Я хочу новый тензор, в котором каждая матрица была сдвинута столбец, например, так

0 0 0
0 0 0
1 1 1

1 1 1
0 0 0
0 0 0

До сих пор я был в состоянии получить потенциальные индексы отображения

import numpy as np

# Input
A = np.array([[0, 0, 0],
              [1, 1, 1],
              [0, 0, 0],
              [0, 0, 0]])
B = np.array([[0, 0, 0],
              [1, 1, 1],
              [0, 0, 0],
              [0, 0, 0]])
AB = np.array([A, B])

# Shifting amount
r = np.array([-1, 1])

d1, d2, d3 = np.ogrid[:AB.shape[0], :AB.shape[1], :AB.shape[2]]

r[r < 0] += AB.shape[1]
r = np.array([r, ]*AB.shape[1]).transpose()
r = r[:, np.newaxis]

# New column indices?
d2 = d2 - r
d2[d2 < 0] += AB.shape[1]

result = AB[d2]
print(result)

Но я получаю эту ошибку:

~/Work/ethz/iml/task2 $ python test.py
Traceback (most recent call last):
  File "test.py", line 27, in <module>
    result = AB[d2]
IndexError: index 2 is out of bounds for axis 0 with size 2

Это то, что d2 выглядит так:

[[[1 1 1 1]
  [2 2 2 2]
  [3 3 3 3]
  [0 0 0 0]]

 [[3 3 3 3]
  [0 0 0 0]
  [1 1 1 1]
  [2 2 2 2]]]

Answer 1

Подход № 1

Адаптация strided-based решения из тех же связанных вопросов и ответов для производительности -

from skimage.util.shape import view_as_windows

def roll_along_second_axis_3dar(a, r):
    r = np.asarray(r)
    a_ext = np.concatenate((a,a[:,:-1,:]),axis=1)
    n = a.shape[1]
    idx = (n-r)%n
    w = view_as_windows(a_ext,(1,n,1))[...,0,:,0]
    return w[np.arange(len(idx)),idx].swapaxes(1,2)

Пробный прогон -

In [11]: a
Out[11]: 
array([[[44, 47, 64],
        [67, 67,  9],
        [83, 21, 36],
        [87, 70, 88]],

       [[88, 12, 58],
        [65, 39, 87],
        [46, 88, 81],
        [37, 25, 77]]])

In [12]: roll_along_second_axis_3dar(a, r=[-1,1])
Out[12]: 
array([[[67, 67,  9],
        [83, 21, 36],
        [87, 70, 88],
        [44, 47, 64]],

       [[37, 25, 77],
        [88, 12, 58],
        [65, 39, 87],
        [46, 88, 81]]])

Подход № 2

Идя с вашей попытки, кажется, вы были достаточно близко. Мы могли бы получить к окончательному выводу несколько изменений / исправлений -

d1, d2, d3 = np.ogrid[:AB.shape[0], :AB.shape[1], :AB.shape[2]]
r[r < 0] += AB.shape[1]
D2 = ((d2 - r).transpose(2,1,0))%AB.shape[1]
out = AB[d1,D2,d3]

Answer 2

Вот подход индексации с take_along_axis:

idx = (-r[:,None] + np.arange(AB.shape[1])[None]) % AB.shape[1]
res = np.take_along_axis(AB, idx[...,None], 1)

Answer 3

Так как в любом случае для перемещения данных потребуется переместить данные в новый массив, вы можете создать список индексов, с которыми будет сопоставляться каждая строка. Мы можем переставить элементы A вдоль оси, выбранной вами в качестве оси «строки», сохранив все остальные оси на месте.

def roll_axis_elements(A, r, raxis=-2, caxis=-1):
    A = np.array(A, copy=False, subok=True)
    *index, = np.indices(A.shape, sparse=True)
    shape = list(A.shape)
    shape[raxis] = shape[caxis] = 1
    r = np.reshape(r, shape)
    index[raxis] = (index[raxis] - r) % A.shape[raxis]
    return A[tuple(index)]

Это будет работать для любой входной формы A , raxis - это ось «строки», которую вы хотите переставить, а caxis - столбцы матрицы, которые перемещаются как единое целое. Вы можете тривиально расширить это до нескольких осей, установив все их в 1, прежде чем изменить r, что также должно исключить их из его формы. Все остальные оси остаются как есть. Это в основном обобщенная версия ответа Дивакара .

Пример прогона:

>>> A = [[[0, 0, 0], [1, 1, 1], [0, 0, 0], [2, 2, 2]],
         [[0, 0, 0], [1, 1, 1], [0, 0, 0], [2, 2, 2]]]
>>> r = [1, -2]
>>> roll_axis_elements(A, r)
array([[[2, 2, 2],
        [0, 0, 0],
        [1, 1, 1],
        [0, 0, 0]],

       [[0, 0, 0],
        [2, 2, 2],
        [0, 0, 0],
        [1, 1, 1]]])