如何对Numpy的2D坐标数组进行仅轴方向的环绕填充？

首先替换低效的坐标生成方式，用广播机制高效创建2D坐标数组：

import numpy as np

map_height = 5
map_width = 5

# 高效生成(5,5,2)的坐标数组
rows = np.arange(map_height)[:, np.newaxis]
cols = np.arange(map_width)[np.newaxis, :]
a_2d_coords = np.stack([rows, cols], axis=-1)

接下来解决环绕填充的核心问题：np.pad默认会对所有维度应用填充，只需通过pad_width参数指定仅对前两个轴（行、列轴）填充1，第三轴（坐标通道）不填充即可：

# 为每个轴单独设置填充宽度：前两轴前后各填1，第三轴不填充
a_2d_coords_wrapped = np.pad(
    a_2d_coords,
    pad_width=((1, 1), (1, 1), (0, 0)),
    mode='wrap'
)

# 验证结果shape：(7, 7, 2)，符合预期
print(a_2d_coords_wrapped.shape)

# 验证x轴环绕：原数组(1,4)的右侧坐标为(1,0)
print(a_2d_coords_wrapped[2, 6])  # 输出 [1 0]
# 验证y轴环绕：原数组(0,0)的上方坐标为(4,0)
print(a_2d_coords_wrapped[0, 1])  # 输出 [4 0]

关键说明 #

• pad_width的格式为((轴0前填充数, 轴0后填充数), (轴1前填充数, 轴1后填充数), ...)，针对每个维度单独配置，第三轴设(0,0)就不会被填充。
• 用广播+np.stack生成坐标数组，比双重循环效率提升显著，尤其适合大规模数组场景。

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