Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Wrap slice around edges of a 2D array in numpy

Tags:

python

numpy

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.

like image 603
George Osterweil Avatar asked Jan 28 '14 03:01

George Osterweil


People also ask

How do you slice a 2D NP array?

Slice Two-dimensional Numpy Arrays To slice elements from two-dimensional arrays, you need to specify both a row index and a column index as [row_index, column_index] . For example, you can use the index [1,2] to query the element at the second row, third column in precip_2002_2013 .

Can you slice NumPy arrays?

You can slice a numpy array is a similar way to slicing a list - except you can do it in more than one dimension. As with indexing, the array you get back when you index or slice a numpy array is a view of the original array. It is the same data, just accessed in a different order.


4 Answers

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.

like image 107
TimCera Avatar answered Oct 17 '22 17:10

TimCera


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.

like image 27
George Osterweil Avatar answered Oct 17 '22 18:10

George Osterweil


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]])
like image 27
tom10 Avatar answered Oct 17 '22 16:10

tom10


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]]
like image 29
Jev Avatar answered Oct 17 '22 16:10

Jev