Python NumPy实现二维数组x/y方向循环平移生成目标数组

• 现有存储坐标信息的4×4二维NumPy数组：

x = np.array([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]])

• 将其展平为一维数组的结果为：x = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15])
• 需求为对该4×4晶格数组分别沿x方向、y方向施加循环平移变换：
  • 沿x方向平移1位时，展平结果需为[1, 2, 3, 0, 5, 6, 7, 4, 9, 10, 11, 8, 13, 14, 15, 12]
  • 沿x方向平移2位时，展平结果需为[2, 3, 0, 1, 6, 7, 4, 5, 10, 11, 8, 9, 14, 15, 12, 13]
• 最终需要生成格式如下的16×16展平二维数组y：

y = np.array([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
              [1, 2, 3, 0, 5, 6, 7, 4, 9, 10, 11, 8, 13, 14, 15, 12],
              [2, 3, 0, 1, 6, 7, 4, 5, 10, 11, 8, 9, 14, 15, 12, 13],
              [3, 0, 1, 2, 7, 4, 5, 6, 11, 8, 9, 10, 15, 12, 13, 14],
              [4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 0, 1, 2, 3],
              [5, 6, 7, 4, 9, 10, 11, 8, 13, 14, 15, 12, 1, 2, 3, 0],
              [6, 7, 4, 5, 10, 11, 8, 9, 14, 15, 12, 13, 2, 3, 0, 1],
              [7, 4, 5, 6, 11, 8, 9, 10, 15, 12, 13, 14, 3, 0, 1, 2],
              [8, 9, 10, 11, 12, 13, 14, 15, 0, 1, 2, 3, 4, 5, 6, 7],
              [9, 10, 11, 8, 13, 14, 15, 12, 1, 2, 3, 0, 5, 6, 7, 4],
              [10, 11, 8, 9, 14, 15, 12, 13, 2, 3, 0, 1, 6, 7, 4, 5],
              [11, 8, 9, 10, 15, 12, 13, 14, 3, 0, 1, 2, 7, 4, 5, 6],
              [12, 13, 14, 15, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
              [13, 14, 15, 12, 1, 2, 3, 0, 5, 6, 7, 4, 9, 10, 11, 8],
              [14, 15, 12, 13, 2, 3, 0, 1, 6, 7, 4, 5, 10, 11, 8, 9],
              [15, 12, 13, 14, 3, 0, 1, 2, 7, 4, 5, 6, 11, 8, 9, 10]])

• 现有错误实现：嵌套调用np.roll的代码y = np.roll(np.roll(x, -1), -1)无法得到预期结果。

之前的np.roll调用未指定axis参数，默认会对展平后的整个数组做平移，不符合沿晶格x/y方向逐行、逐块平移的要求。
正确实现逻辑：遍历y方向0-3共4个平移步长，再遍历x方向0-3共4个平移步长；每次平移时先对原始二维晶格沿y轴（行方向）做对应步长的循环平移，再沿x轴（列方向）做对应步长的循环平移，将每次平移后的结果展平，按顺序堆叠即可得到16×16的目标数组。

实现代码：

import numpy as np

# 原始4×4晶格数组
x = np.array([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]])
grid_size = x.shape[0]
result = []

for dy in range(grid_size):
    for dx in range(grid_size):
        # 按y、x方向依次做循环平移
        shifted_arr = np.roll(np.roll(x, dy, axis=0), dx, axis=1)
        result.append(shifted_arr.flatten())

y = np.array(result)

运行后得到的y与给出的目标格式完全一致。如果需要向前平移而非向后循环移位，将np.roll的步长参数改为-dy、-dx即可。

内容的提问来源于stack exchange，提问作者relaxon