wrapping around slices in Python / numpy

Question

I have a numpy array, and I want to get the "neighbourhood" of the i'th point. Usually the arrays I'm using are two-dimensional, but the following 1D example illustrates what I'm looking for. If

A = numpy.array([0,10,20,30,40,50,60,70,80,90])

Then the (size 5) neighbourhood of element 4 is [20,30,40,50,60], and this can easily be obtained by doing A[i-2:i+3].

However, I also need the neighbourhoods to "wrap around" the edges of the array, so that the neighbourhood of the element 0 is [80,90,0,10,20] and the neighbourhood of the element 9 is [70,80,90,0,10]. I can't seem to find an elegant way to do this, so I end up having to use some complicated, annoying logic every time this comes up (which is very often for me). In the 2D case the neighbourhood of a point would be a rectangular array.

So my question is, is there a neat way to expres this "wrap-around neighbourhood" operation in numpy? I would prefer something that returns a slice rather than a copy, but readability and speed are the most important considerations.

Answer 1

numpy.take in 'wrap' mode will use your indices modulo the length of the array.

indices = range(i-2,i+3)
neighbourhood = A.take(indices, mode='wrap')

See documentation for details numpy.take