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

you can use argument axis=0 of numpy.take for n-d array.

A = zip(range(0,101,10),range(0,11)) #create 2-d list
A = numpy.array(A)   #create 2-d array  
indices = range(i-2,i+3)
neightbourhood = A.take(indices,axis=0,mode='wrap')

The same axis=0 will work for n*m dimensions...

Answer 2

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

Answer 3

numpy.roll can shift the array such that the entire slice is at the beginning of the array. Then take your slice at the beginning and numpy.roll again to revert the array back to its original position.

# modify array at index i and nearest two
# locations on each side of i, wrapping
# around the edges
A = np.array([0,10,20,30,40,50,60,70,80,90])
i = 9
neighbors = 2
A=np.roll(A, -i+neighbors)
A[:5] += 1
A=np.roll(A, i-neighbors)

array([ 1, 11, 20, 30, 40, 50, 60, 71, 81, 91])

numpy.roll doesn't perform well for me on large arrays however.

Answer 4

I know this question is old, but should mention scipy.ndimage.filter.generic_filter.

It has a mode='wrap' option, plus it handles the application of the neighbor function.

import scipy.ndimage as nd

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

Say you have a neighbor function:

def nbf(arr):
    return sum(arr)

To apply the neighbor function to every 5, with wrapped values at the edges:

C = nd.generic_filter(A, nbf, 5, mode='wrap')

print(C)
[200 150 100 150 200 250 300 350 300 250]

Answer 5

Note: For cases where your neighbors do not require wrapping, numpy.take is slower than simply taking a slice A[i-2:i+3]. You may want to wrap your neighbors function with some conditional statements:

def neighbors(a,i,n):
    N = a.shape[0] 
    if i - n < 0 and i + n > 0:
        indices = range(i-n,i+n+1)
        nbrs = a.take(indices, mode='wrap')
    elif i-n < N - 1 and i+n > N - 1:
        indices = range(i-n,i+n+1)
        nbrs = a.take(indices, mode='wrap')
    else:
        nbrs = a[i-n:i+n+1]
    return nbrs

If you find yourself taking neighbors while iterating through an array, like in a centered moving average, you'll find that this requires less time, especially for longer arrays:

Here is the moving average function I used:

def moving_average(a,n=1):
    N = a.shape[0] 
    ma = np.empty(N)
    for i in range(N):
        if n*2+1 > N:
            ma[i] = a.mean()
        else: 
            ma[i] = neighbors(a,i,n).mean()
    return ma

I'm sure these function can be improved further. I'm open to suggestions.

Answer 6

You can use the np.pad routine like this:

A = np.array([0,10,20,30,40,50,60,70,80,90])
A = np.pad(A, 2, 'wrap')
print(A)
[80, 90,  0, 10, 20, 30, 40, 50, 60, 70, 80, 90,  0, 10]

Say you have a neighbor function:

def nbf(arr):
    return sum(arr)

To apply the neighbor function to every 5 you need to be careful about your start and end indices (in the range(...) command) and the relative slice you take from A.

B = [nbf(A[i-2:i+3]) for i in range(2,12)]
print(B)
[200, 150, 100, 150, 200, 250, 300, 350, 300, 250]