Python，规则网格上的邻居答案

【问题标题】：Python, neighbors on a regular gridPython，规则网格上的邻居
【发布时间】：2013-03-29 01:15:04
【问题描述】：

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

如果给每个单元格和索引分配如下定义，问题就很简单了：

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]

你认为有最快的方法来实现这个算法吗？

【问题讨论】：

我认为你可以通过使用 NumPy 花式索引来加速它。你能发布更多创建idx, idy, cells, idx_cells,x的代码吗？
你也可以使用 scipy 的 cKDTree。 docs.scipy.org/doc/scipy/reference/generated/…
谢谢，我已经尝试过使用 KDTree，但是对于我的数据集来说它相当慢，而且对于网格来说也不方便。不幸的是，我无法发布更多代码，因为我的数据集很大。

标签： python numpy

【解决方案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 关键字参数，您可以使用它来处理格子边缘的项目，请参阅文档。

【讨论】：