Python 多维网格中周期边界条件的实现

我尝试实现一种算法，在具有周期性边界条件的N维网格中查找点的邻域

例如，我有一个长度为3（3x3x3）的立方体三维网格，所以我有27个点。每个点都通过

numpy.ravel\u multi\u index

（见图）分配到索引。 现在我想找到一个任意点的邻居。我的代码适用于内部点：

def nneighbour(index , dim, Length):
    neighbour = np.zeros(dim*2) 
    up_down = np.empty((dim*2,),int)
    up_down[::2] = 1
    up_down[1::2] = -1 # array for left right -1 = left, +1 = right
    D = np.repeat(np.arange(0, dim, 1), 2) # neighbour dimension : first 2 elements in 1d, next 2 elements in 2d and so on
    neighbour = index + up_down*np.power(Length, D)
    return neighbour

nneighbour（13,3,3）

返回例如

[14 12 16 10 22 4]

。前两个点是第一维的邻域，后两个点是第二维的邻域，依此类推。 但我不知道如何实现边点的周期边界条件。对于

index=10

，它应该是

[11 9 13 16 19 1]

而不是

[11 9 13 7 19 1]

（因为我们在第二层中跨越了底部边界）

---

.

最简单的方法是解开-做邻居-解开：

def nneighbour(index,dim,length):
    index = np.asarray(index)
    shp = dim*(length,)
    neighbours = np.empty((*index.shape,dim,dim,2),int)
    neighbours.T[...] = np.unravel_index(index,shp)
    ddiag = np.einsum('...iij->...ij',neighbours)
    ddiag += (-1,1)
    ddiag %= length
    return np.ravel_multi_index(np.moveaxis(neighbours,-3,0),shp).ravel()

例如：

nneighbour(10,3,3)
# array([ 1, 19, 16, 13,  9, 11])

---

如果我有多个输入索引，是否有一种方法可以通过一次运行获得所有索引的邻居？