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.

note that it is fundamentally impossible to get a view of such a subarray in numpy; the subarray cannot be expressed using a single stride for each axis — Eelco Hoogendoorn, Jan 15 '14 at 08:23

Henrik · Answer 1 · 2013-07-19T13:53:44.097

47

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

edited Jul 19 '13 at 13:53

answered Jul 19 '13 at 06:45

Henrik

4,254
15
28

6

Beat me by 10 seconds. Hmmph! Note that you can also access `take` as an array method, i.e. `A.take(indices, mode='wrap')`. – DSM Jul 19 '13 at 06:46
Thanks @DSM, edited as it makes more sense to use the instance's own method in this case. – Henrik Jul 19 '13 at 06:55
Ah. I forgot to mention that my array is two dimensional - I just used a 1D array to make the problem clearer. – N. Virgo Jul 19 '13 at 12:18
1

In my current case I want to do something like `A[column, take(row_indices,mode='wrap')]`, if you see what I mean, but I also often need to get a full n-by-m neighbourhood of a point. (+1 for a good answer to the question I asked, even if it wasn't quite the question I intended to ask :). I didn't know about the `take` method.) – N. Virgo Jul 19 '13 at 12:21
@Nathaniel I think u can use the same answer, because here u r dealing with index of array to get neighbours, not the element. AFAIT it should work, irrespective of ur array dimensions. – tailor_raj Jul 21 '13 at 06:21
@tailor_raj can you give an example of using `take` with a 2D array? There is nothing about it in the documentation. I tried the most obvious method (pass a list of tuples as the `indices` argument) but that doesn't work as hoped. – N. Virgo Jul 21 '13 at 06:29
@Nathaniel See my answer. If it fulfills ur purpose. – tailor_raj Jul 21 '13 at 08:43
Be warned that this allocates a new array, not provides a view into the original one like normal slices. So, unlike slices, this can't be used to mutate the original array! – ivan_pozdeev Oct 30 '16 at 14:50

tailor_raj · Answer 2 · 2014-01-15T06:14:19.147

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

edited Jan 15 '14 at 06:14

answered Jul 21 '13 at 07:52

tailor_raj

1,037
2
9
19

score 6 · Answer 3 · answered Mar 25 '16 at 14:49

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.

score 4 · Answer 4 · edited Sep 23 '17 at 05:16

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]

score 3 · Answer 5 · answered Oct 25 '14 at 19:55

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]

score 3 · Answer 6 · answered Jun 24 '15 at 20:57

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.

It should be noted that numpy.roll will allow shifting along multiple dimensions if you pass a tuple for shift parameter and a tuple for axis parameter. — bfris, Sep 08 '21 at 19:39

score 1 · Answer 7 · answered Oct 08 '20 at 22:48

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

wrapping around slices in Python / numpy

7 Answers7

Linked