python - Wrap slice around edges of a 2D array in numpy

Question

Suppose I am working with numpy in Python and I have a two-dimensional array of arbitrary size. For convenience, let's say I have a 5 x 5 array. The specific numbers are not particularly important to my question; they're just an example.

a = numpy.arrange(25).reshape(5,5)

This yields:

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

Now, let's say I wanted to take a 2D slice of this array. In normal conditions, this would be easy. To get the cells immediately adjacent to 2,2 I would simply use a[1:4,1,4] which would yield the expected

[[6, 7,   8 ],
 [11, 12, 13],
 [16, 17, 18]]

But what if I want to take a slice that wraps around the edges of the array? For example a[-1:2,-1:2] would yield:

[24, 20, 21],
[4, 0,  1 ],
[9, 5,  6 ] 

This would be useful in several situations where the edges don't matter, for example game graphics that wrap around a screen. I realize this can be done with a lot of if statements and bounds-checking, but I was wondering if there was a cleaner, more idiomatic way to accomplish this.

Looking around, I have found several answers such as this: https://stackoverflow.com/questions/17739543/wrapping-around-slices-in-python-numpy that work for 1-dimensional arrays, but I have yet to figure out how to apply this logic to a 2D slice.

So essentially, the question is: how do I take a 2D slice of a 2D array in numpy that wraps around the edges of the array?

Thank you in advance to anyone who can help.

Answer 1

This will work with numpy >= 1.7.

a = np.arange(25).reshape(5,5)

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

The pad routine has a 'wrap' method...

b = np.pad(a, 1, mode='wrap')

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

Depending on the situation you may have to add 1 to each term of any slice in order to account for the padding around b.