0
votes

I try to implement an algorithm that finds the neighbours of a point in N-dimensional grids with periodic boundary conditions.

For example I have a cubic 3D grid with Length = 3 (3x3x3), so I have 27 points. Every point gets assigned to an index via numpy.ravel_multi_index(see picture). Now I want to find the neighbours for an arbitrary point. My code works for inner points:

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) returns for example [14 12 16 10 22 4]. First two points are neighbors in the first dimension, next two for 2nd dimension and so on. But I have no idea how to implement periodic boundary conditions for edge points. For index = 10 it should be [11 9 13 16 19 1] not [11 9 13 7 19 1] (as we cross the bottom boundary in the second level).

Flattened 3x3x3 grid.

1

1 Answers

1
votes

The easiest way is to unravel - do the neighbours - ravel:

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

Example:

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