查找多维数组中元素的所有 "neighbors",环绕边界

Find all "neighbors" for element in multidimensional array, wrapping around the boundaries

假设我使用 2D/3D numpy 数组来模拟由单元格组成的盒子。单元格由从 0 开始递增的数字标记。到这里为止,这可以通过

完成
box = np.arange(np.prod(n_cells))
box = box.reshape(n_cells)

其中 n_cells 是一个 np.array,它存储框在每个维度上应具有的单元格数,基本上是按单元格计算的框的尺寸。

现在,对我来说最困难的部分是找到一种智能方法来为每个单元格找到邻居。我的目标是建立一个数组或列表,其中包含每个单元格的邻居 - 理想情况下让我们看一个小的二维示例。

0   1   2   3
4   5   6   7
8   9   10 11
12  13  14 15 

这里,没有周期性边界条件,邻居是

0: 1, 4, 5
1: 0, 2, 4, 5, 6
2: 1, 5, 6, 7, 3
...

但我希望它具有周期性边界条件,这样

0: 1, 4, 5, 3, 7, 12, 13, 15
1: 0, 2, 4, 5, 6, 12, 13, 14

这样,每个元素在 2D 中都有 8 个相邻元素。 理想情况下,我希望能够为任何维度创建这样的 list/array,但如果没有通用解决方案,我对 2D/3D 特别感兴趣。

另一个问题,我必须解决的是这样,我对所有的对计数两次,即 0 是 1 的邻居,1 是 0 的邻居。这是我需要摆脱的东西也是,但这不是主要问题。

这可以通过 numpy.roll 来完成,其中 "rolls" 沿给定轴的数组,具有您想要的那种环绕。例如,按 (-1, -1) 滚动会将所有内容向左和向上移动,因此数组变为

  [[ 5,  6,  7,  4],
   [ 9, 10, 11,  8],
   [13, 14, 15, 12],
   [ 1,  2,  3,  0]]

这样,我们就找到了每个点的东南邻居。剩下的就是展平这个列表 (ravel),对 9 个偏移中的每一个重复这个过程(包括 (0, 0),这意味着数字本身),然后堆叠结果。该解决方案适用于数组的任意维度 b:

dim = len(b.shape)       # number of dimensions
offsets = [0, -1, 1]     # offsets, 0 first so the original entry is first 
columns = []
for shift in itertools.product(offsets, repeat=dim):   # equivalent to dim nested loops over offsets
    columns.append(np.roll(b, shift, np.arange(dim)).ravel())
neighbors = np.stack(columns, axis=-1)

输出(neighbors的值):

  [[ 0,  1,  3,  4,  5,  7, 12, 13, 15],
   [ 1,  2,  0,  5,  6,  4, 13, 14, 12],
   [ 2,  3,  1,  6,  7,  5, 14, 15, 13],
   [ 3,  0,  2,  7,  4,  6, 15, 12, 14],
   [ 4,  5,  7,  8,  9, 11,  0,  1,  3],
   [ 5,  6,  4,  9, 10,  8,  1,  2,  0],
   [ 6,  7,  5, 10, 11,  9,  2,  3,  1],
   [ 7,  4,  6, 11,  8, 10,  3,  0,  2],
   [ 8,  9, 11, 12, 13, 15,  4,  5,  7],
   [ 9, 10,  8, 13, 14, 12,  5,  6,  4],
   [10, 11,  9, 14, 15, 13,  6,  7,  5],
   [11,  8, 10, 15, 12, 14,  7,  4,  6],
   [12, 13, 15,  0,  1,  3,  8,  9, 11],
   [13, 14, 12,  1,  2,  0,  9, 10,  8],
   [14, 15, 13,  2,  3,  1, 10, 11,  9],
   [15, 12, 14,  3,  0,  2, 11,  8, 10]]

在每一行中,第一个条目是原始数字,其他是它的邻居。

要让每个条目-邻居对仅列出一次,您可以屏蔽冗余条目,例如使用 NaN:

np.where(neighbors >= neighbors[:, [0]], neighbors, np.nan)

  [[  0.,   1.,   3.,   4.,   5.,   7.,  12.,  13.,  15.],
   [  1.,   2.,  nan,   5.,   6.,   4.,  13.,  14.,  12.],
   [  2.,   3.,  nan,   6.,   7.,   5.,  14.,  15.,  13.],
   [  3.,  nan,  nan,   7.,   4.,   6.,  15.,  12.,  14.],
   [  4.,   5.,   7.,   8.,   9.,  11.,  nan,  nan,  nan],
   [  5.,   6.,  nan,   9.,  10.,   8.,  nan,  nan,  nan],
   [  6.,   7.,  nan,  10.,  11.,   9.,  nan,  nan,  nan],
   [  7.,  nan,  nan,  11.,   8.,  10.,  nan,  nan,  nan],
   [  8.,   9.,  11.,  12.,  13.,  15.,  nan,  nan,  nan],
   [  9.,  10.,  nan,  13.,  14.,  12.,  nan,  nan,  nan],
   [ 10.,  11.,  nan,  14.,  15.,  13.,  nan,  nan,  nan],
   [ 11.,  nan,  nan,  15.,  12.,  14.,  nan,  nan,  nan],
   [ 12.,  13.,  15.,  nan,  nan,  nan,  nan,  nan,  nan],
   [ 13.,  14.,  nan,  nan,  nan,  nan,  nan,  nan,  nan],
   [ 14.,  15.,  nan,  nan,  nan,  nan,  nan,  nan,  nan],
   [ 15.,  nan,  nan,  nan,  nan,  nan,  nan,  nan,  nan]])

想法是 neighbors >= neighbors[:, [0]] 只列出那些数字大于单元格本身的人。