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

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.