python：独立地移动张量中的每个矩阵

Question

这个问题对于矩阵Roll rows of a matrix independently非常类似

但我没能适应三维张量

给我一个张量

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

从同一链接Q&；A就表演而言—

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

由于滚动无论如何都需要将数据移动到一个新数组中，所以您可以创建一个索引列表，每个行都将映射到该列表。我们可以沿着您选择为“行”轴的轴排列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是作为一个单元移动的矩阵列。通过在重塑r之前将所有轴都设置为1，可以轻松地将其扩展到多个轴，这也必须将它们从其形状中排除。所有其他轴保持原样。这基本上是Divakar's answer的广义版本

样本运行：

>>> 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]]])

Answer 3

以下是一种使用take_along_axis的索引方法：

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