CODEWITHSUNDEEP
pythonnumpy
N. Virgo
N. Virgo

Reputation: 8439

wrapping around slices in Python / numpy

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.

Upvotes: 41

Views: 23852

Answers (7)

saladi
saladi

Reputation: 3253

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

Upvotes: 1

TimCera
TimCera

Reputation: 592

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]

Upvotes: 4

Neal Kruis
Neal Kruis

Reputation: 2287

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:

enter image description here

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.

Upvotes: 8

Joseph Sheedy
Joseph Sheedy

Reputation: 6726

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.

Upvotes: 3

TimCera
TimCera

Reputation: 592

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]

Upvotes: 3

tailor_raj
tailor_raj

Reputation: 1057

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

Upvotes: 6

Henrik
Henrik

Reputation: 4314

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

Upvotes: 51

Related Questions