Python, neighbors on a regular grid

Question

Let's suppose I have a set of 2D coordinates that represent the centers of cells of a 2D regular mesh. I would like to find, for each cell in the grid, the two closest neighbors in each direction. The problem is quite straightforward if one assigns to each cell and index defined as follows: idx_cell = idx+N*idy where N is the total number of cells in the grid, idx=x/dx and idy=y/dx, with x and y being the x-coordinate and the y-coordinate of a cell and dx its size. For example, the neighboring cells for a cell with idx_cell=5 are the cells with idx_cell equal to 4,6 (for the x-axis) and 5+N,5-N (for the y-axis). The problem that I have is that my implementation of the algorithm is quite slow for large (N>1e6) data sets. For instance, to get the neighbors of the x-axis I do `[x[(idx_cell==idx_cell[i]-1)|(idx_cell==idx_cell[i]+1)] for i in cells]` Do you think there's a fastest way to implement this algorithm?

Accepted Answer

You are basically reinventing the indexing scheme of a multidimensional array. It is relatively easy to code, but you can use the two functions `unravel_index` and `ravel_multi_index` to your advantage here.

If your grid is of `M` rows and `N` columns, to get the `idx` and `idy` of a single item you could do:

```
>>> M, N = 12, 10
>>> np.unravel_index(4, dims=(M, N))
(0, 4)
```

This also works if, instead of a single index, you provide an array of indices:

```
>>> np.unravel_index([15, 28, 32, 97], dims=(M, N))
(array([1, 2, 3, 9], dtype=int64), array([5, 8, 2, 7], dtype=int64))
```

So if `cells` has the indices of several cells you want to find neighbors to:

```
>>> cells = np.array([15, 28, 32, 44, 87])
```

You can get their neighbors as:

```
>>> idy, idx = np.unravel_index(cells, dims=(M, N))
>>> neigh_idx = np.vstack((idx-1, idx+1, idx, idx))
>>> neigh_idy = np.vstack((idy, idy, idy-1, idy+1))
>>> np.ravel_multi_index((neigh_idy, neigh_idx), dims=(M,N))
array([[14, 27, 31, 43, 86],
       [16, 29, 33, 45, 88],
       [ 5, 18, 22, 34, 77],
       [25, 38, 42, 54, 97]], dtype=int64)
```

Or, if you prefer it like that:

```
>>> np.ravel_multi_index((neigh_idy, neigh_idx), dims=(M,N)).T
array([[14, 16,  5, 25],
       [27, 29, 18, 38],
       [31, 33, 22, 42],
       [43, 45, 34, 54],
       [86, 88, 77, 97]], dtype=int64)
```

The nicest thing about going this way is that `ravel_multi_index` has a `mode` keyword argument you can use to handle items on the edges of your lattice, see the docs.

Python, neighbors on a regular grid

Solution

Problem