用周期条件确定二维网格中的相邻单元

Question

假设我们有以下二维网络，我们用整数标记其单元索引：

20  21  22  23  24
15  16  17  18  19
10  11  12  13  14
5   6   7   8   9
0   1   2   3   4

我想要的是一个函数，它接收一个单元格索引（cell）和一个轴上的单元格数（本例中为n=5），并返回一个包含9个相邻单元格（包括单元格本身）的数组，同时考虑全局框的周期性

我向您展示了我尝试过的“几乎”有效的方法：

def celdas_vecinas(cell,n):

    return np.mod(cell + np.array([0, -n-1, -n, -n+1, -1, 1, n-1, n, n+1], dtype=np.int64), n**2)

其中我输入了np.mod以反映周期条件。关键是该函数仅对某些值表现良好

>>> celdas_vecinas(1,5) 
array([ 1, 20, 21, 22,  0,  2,  5,  6,  7]) # right!

>>> celdas_vecinas(21,5)
array([21, 15, 16, 17, 20, 22,  0,  1,  2]) # right!

但如果我在角落中输入其中一个单元格的索引，则会发生以下情况：

>>> celdas_vecinas(0,5)
array([ 0, 19, 20, 21, 24,  1,  4,  5,  6]) # should be 9 instead of 19

例如，对于单元格=5也会失败

有人知道我如何实现这个功能吗？当单元索引不接触任何边界时，它很容易实现，但我不知道如何包含周期性效应，尽管我猜它一定与np.mod函数有关

Answer 1

Numpy实现可以利用numpy.argwhere检索值索引，使用numpy.ix_创建索引网格，最后应用numpy.narray.ravel方法将数组展开：

import numpy as np

n = 5
grid = np.arange(n**2).reshape(n,n)[::-1]

def celdas_vecinas_np(grid, v, n):

    x, y = np.argwhere(grid == v)[0]
    idx = np.arange(x-1, x+2) %n
    idy = np.arange(y-1, y+2) %n

    return grid[np.ix_(idx, idy)].ravel()

celdas_vecinas_np(grid, 24, n)
array([ 3,  4,  0, 23, 24, 20, 18, 19, 15])

另一方面，对于Numba实现我们不能使用numpy.argwhere，但我们可以使用numpy.where来获取索引。一旦我们这样做，就只需要在正确的范围内循环，即：

from numba import njit

@njit
def celdas_vecinas_numba(grid, v, n):

    x, y = np.where(grid == v)
    x, y = x[0], y[0]

    result = []
    for ix in range(x-1, x+2):
        for iy in range(y-1, y+2):
            result.append(grid[ix%n, iy%n])

    return result

celdas_vecinas_numba(grid, 24, n)
[3, 4, 0, 23, 24, 20, 18, 19, 15]

性能比较使用如此小的网格，numba在我的本地机器上的运行速度已经快了约20倍：

%timeit celdas_vecinas_np(grid, 24, 5)
38 µs ± 1.62 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit celdas_vecinas_numba(grid, 24, n)
1.81 µs ± 93.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)