在Python / numpy中围绕切片

Question

我有一个numpy数组，我想得到第i个点的“邻域”。通常我使用的数组是二维的，但下面的1D示例说明了我正在寻找的东西。如果

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

那么元素4的(大小为5)邻域是[20,30,40,50,60]，这可以通过执行A[i-2:i+3]很容易地获得。

但是，我还需要“环绕”数组边缘的邻域，以便元素0的邻域是[80,90,0,10,20]，元素9的邻域是[70,80,90,0,10]。我似乎找不到一种优雅的方法来做到这一点，所以每次出现这种情况时，我都不得不使用一些复杂的、令人讨厌的逻辑(这对我来说是非常常见的)。在2D情况下，一个点的邻域将是一个矩形阵列。

所以我的问题是，有没有一种巧妙的方法来表达numpy中的“环绕邻居”操作？我更喜欢返回切片而不是副本的东西，但可读性和速度是最重要的考虑因素。

Answer 1

'wrap'模式下的numpy.take将使用以数组长度为模的索引。

indices = range(i-2,i+3)
neighbourhood = A.take(indices, mode='wrap')

有关详细信息numpy.take，请参阅文档

Answer 2

对于n-d数组，可以使用numpy.take的参数axis=0。

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

同样的axis=0将适用于n*m维...

Answer 3

注意：对于您的邻居不需要包装的情况，numpy.take比简单地获取切片A[i-2:i+3]慢。你可能想用一些条件语句来包装你的邻居函数：

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

如果你发现自己在迭代数组时需要邻居，比如在中心移动平均中，你会发现这需要更少的时间，特别是对于更长的数组：

​

​

下面是我使用的移动平均函数：

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

我相信这些功能还可以进一步改进。我乐于接受大家的建议。