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.

Answer 2

After playing around with various methods for a while, I just came to a fairly simple solution that works using ndarray.take. Using the example I provided in the question:

a.take(range(-1,2),mode='wrap', axis=0).take(range(-1,2),mode='wrap',axis=1)

Provides the desired output of

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

It turns out to be a lot simpler than I thought it would be. This solution also works if you reverse the two axes.

This is similar to the previous answers I've seen using take, but I haven't seen anyone explain how it'd be used with a 2D array before, so I'm posting this in the hopes it helps someone with the same question in the future.

Answer 3

You can also use roll, to roll the array and then take your slice:

b = np.roll(np.roll(a, 1, axis=0), 1, axis=1)[:3,:3]

gives

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

Answer 4

I had a similar challenge working with wrap-around indexing, only in my case I needed to set values in the original matrix. I've solved this by 'fancy indexing' and making use of meshgrid function:

A = arange(25).reshape((5,5)) # destinatoin matrix
print 'A:\n',A

k =-1* np.arange(9).reshape(3,3)# test kernel, all negative
print 'Kernel:\n', k
ix,iy = np.meshgrid(arange(3),arange(3)) # create x and y basis indices

pos = (0,-1) # insertion position

# create insertion indices
x = (ix+pos[0]) % A.shape[0]
y = (iy+pos[1]) % A.shape[1]
A[x,y] = k # set values
print 'Result:\n',A

The output:

A:
[[ 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]]
Kernel:
[[ 0 -1 -2]
 [-3 -4 -5]
 [-6 -7 -8]]
Result:
[[-3 -6  2  3  0]
 [-4 -7  7  8 -1]
 [-5 -8 12 13 -2]
 [15 16 17 18 19]
 [20 21 22 23 24]]

Answer 5

As I mentioned in the comments, there is a good answer at How do I select a window from a numpy array with periodic boundary conditions?

Here is another simple way to do this

# First some setup
import numpy as np
A = np.arange(25).reshape((5, 5))
m, n = A.shape

and then

A[np.arange(i-1, i+2)%m].reshape((3, -1))[:,np.arange(j-1, j+2)%n]

It is somewhat harder to obtain something that you can assign to. Here is a somewhat slower version. In order to get a similar slice of values I would have to do

A.flat[np.array([np.arange(j-1,j+2)%n+a*n for a in xrange(i-1, i+2)]).ravel()].reshape((3,3))

In order to assign to this I would have to avoid the call to reshape and work directly with the flattened version returned by the fancy indexing. Here is an example:

n = 7
A = np.zeros((n, n))
for i in xrange(n-2, 0, -1):
    A.flat[np.array([np.arange(i-1,i+2)%n+a*n for a in xrange(i-1, i+2)]).ravel()] = i+1
print A

which returns

[[ 2.  2.  2.  0.  0.  0.  0.]
 [ 2.  2.  2.  3.  0.  0.  0.]
 [ 2.  2.  2.  3.  4.  0.  0.]
 [ 0.  3.  3.  3.  4.  5.  0.]
 [ 0.  0.  4.  4.  4.  5.  6.]
 [ 0.  0.  0.  5.  5.  5.  6.]
 [ 0.  0.  0.  0.  6.  6.  6.]]