Обернуть 2D-координаты массива numpy

Прежде всего, вы можете упростить расчет 2D-координат следующим образом:

a_2d_coords = np.moveaxis(np.mgrid[:map_height, :map_width], 0, -1)

Давайте проверим, дает ли это тот же результат:

import numpy as np

# For the check, use different height and width values to detect swaps
map_height, map_width = 3, 7

# Original implementation for coords
a_2d_coords_given = np.zeros((map_height, map_width, 2), dtype = np.int16)

for row_idx in range(map_height):
    for col_idx in range(map_width):
        a_2d_coords_given[row_idx, col_idx][0] = row_idx
        a_2d_coords_given[row_idx, col_idx][1] = col_idx

# Proposed implementation for coords
a_2d_coords = np.moveaxis(np.mgrid[:map_height, :map_width], 0, -1)
# Equality check
assert np.all(a_2d_coords == a_2d_coords_given)

Во-вторых, если вы заметите, что аргумент pad_widthnp.pad также может принимать значения, специфичные для оси, то вы можете добиться желаемого заполнения координат следующим образом:

a_2d_coords_wrapped = np.pad(a_2d_coords, mode = "wrap",
                             pad_width=((1, 1), (1, 1), (0, 0)))

Теперь мы дополняем ось 0 и ось 1 одним ведущим и одним конечным значением ((1, 1)), но оставляем ось 2 в покое ((0, 0)).

Итого имеем:

import numpy as np

map_height, map_width = 5, 5
a = np.arange(map_height * map_width).reshape(map_height, map_width)

a_2d_coords = np.moveaxis(np.mgrid[:map_height, :map_width], 0, -1)

a_2d_coords_wrapped = np.pad(a_2d_coords, mode = "wrap",
                             pad_width=((1, 1), (1, 1), (0, 0)))
print(a[a_2d_coords_wrapped[..., 0], a_2d_coords_wrapped[..., 1]])
# Prints
# [[24 20 21 22 23 24 20]
#  [ 4  0  1  2  3  4  0]
#  [ 9  5  6  7  8  9  5]
#  [14 10 11 12 13 14 10]
#  [19 15 16 17 18 19 15]
#  [24 20 21 22 23 24 20]
#  [ 4  0  1  2  3  4  0]]

Дальнейшее упрощение

Мы можем упростить код еще на один шаг, предоставив a_2d_coords/a_2d_coords_wrapped как массив 2×В×Ш, а не массив В×Ш×2:

import numpy as np

map_height, map_width = 5, 5
a = np.arange(map_height * map_width).reshape(map_height, map_width)

a_2d_coords = np.mgrid[:map_height, :map_width]

a_2d_coords_wrapped = np.pad(a_2d_coords, mode = "wrap",
                             pad_width=((0, 0), (1, 1), (1, 1)))
print(a[*a_2d_coords_wrapped])
# Prints
# [[24 20 21 22 23 24 20]
#  [ 4  0  1  2  3  4  0]
#  [ 9  5  6  7  8  9  5]
#  [14 10 11 12 13 14 10]
#  [19 15 16 17 18 19 15]
#  [24 20 21 22 23 24 20]
#  [ 4  0  1  2  3  4  0]]

Обратите внимание на * в a[*a_2d_coords_wrapped], который необходим для «распаковки» координат, чтобы для индексации предоставлялись два массива H×W, а не один массив 2×H×W.