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

Answer 1

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]

Answer 2

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

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 5

If you don't have the convenience of using np.take with mode='wrap' (e.g. when using numba), the following works just the same:

A = numpy.array([0,10,20,30,40,50,60,70,80,90])
indices = range(i-2, i+3)
neighbourhood = A.take(indices % len(A))

or

A = numpy.array([0,10,20,30,40,50,60,70,80,90])
indices = range(i-2, i+3)
neighbourhood = A[indices % len(A)]

Answer 6

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...