蟒蛇，在规则网格上的邻居

Question

假设我有一组二维坐标，它们代表二维规则网格的单元格中心。我想找到，对于网格中的每个单元格，在每个方向上的两个最近的邻居。

如果对每个单元格和索引进行如下定义，问题就很简单了:

idx_cell = idx+N*idy

，其中N为网格中单元格的总数，idx=x/dx, idy=y/dx，其中x和y分别为单元格的x坐标和y坐标，dx为单元格的大小。

例如，idx_cell=5的单元格的相邻单元格是idx_cell等于4,6(对于x轴)和5+N,5-N(对于y轴)的单元格。

我的问题是，我的算法的实现是相当慢的大(N>1e6)数据集。

例如，为了得到x轴的邻居，我使用

[x[(idx_cell==idx_cell[i]-1)|(idx_cell==idx_cell[i]+1)] for i in cells]

你认为有没有最快的方法来实现这个算法?

Answer 1

您基本上是在重新发明多维数组的索引方案。它相对容易编写，但是您可以在这里使用unravel_index和ravel_multi_index这两个函数。

如果你的网格是M行和N列，获得idx和idy的一个项目，你可以做:

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

如果不是提供单个索引，而是提供索引数组，也可以这样做:

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

所以如果cells有几个单元格的索引你想找到邻居:

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

你可以得到他们的邻居:

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

或者，如果你喜欢这样:

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

这样做的最好的事情是ravel_multi_index有一个mode关键字参数，您可以使用它来处理格边缘的项，请参阅文档。