I have a (5, 5) 2D numpy array:

map_height = 5
map_width = 5

# Define a 2D np array-
a = np.arange(map_height * map_width).reshape(map_height, map_width)

# a
'''
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24]])
'''

I can wrap this array around on both of its axes using ‘pad()’:

a_wrapped = np.pad(array = a, pad_width = 1, mode = 'wrap')

a_wrapped
'''
array([[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]])
'''

The 2D coordinates of (5, 5) ‘a’ are computed (inefficiently) as:

# 2D coordinates -
# 1st channel/axis = row indices & 2 channel/axis = column indices.
a_2d_coords = 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[row_idx, col_idx][0] = row_idx
        a_2d_coords[row_idx, col_idx][1] = col_idx


# a_2d_coords.shape
# (5, 5, 2)

I want to wrap this 2D coordinates array ‘a_2d_coords’ as well by doing:

a_2d_coords_wrapped = np.pad(array = a_2d_coords, pad_width = 1, mode = 'wrap')

# a_2d_coords_wrapped.shape
# (7, 7, 4)

It also wraps the 3rd axis/dimension which should not be done! The goal is that the coords of a[1, 4] = (1, 4) and its neighbor’s coords to right hand side (RHS) should be a[1, 0] = (1, 0). This is wrapping around the x-axis. Similarly, y-axis 2D coordinates should also be wrapped.

?