将数组元素的邻居广义堆叠成 3D 数组

Question

设置

给定一个 2D 数组，我想创建一个 3D 数组，其中沿第三维的值（即stacked[row, col, :]）是[row, col] 处原始数组的扁平邻居。我想概括这个过程来处理任意（但合理的）搜索半径。

先前的研究

这个question 看起来很有希望，但我不确定我是否可以在没有（几个）for 循环的情况下真正利用它的方法。为简洁起见，我当前的方法（搜索半径为 1）如下面的示例所示。

这个question + answer 也很接近，但我特意寻找一种纯粹使用智能索引来避免循环的解决方案。

我现在拥有的

import numpy as np
np.random.seed(0)

x = np.random.random_integers(0, 10, size=(4, 5))
print(x)  # * highlights the neighbors we'll see later

[[ 5   0   3   3   7]
 [ 9  *3  *5  *2   4]
 [ 7  *6  *8  *8  10]
 [ 1  *6  *7  *7   8]]

# padding the edges 
padded = np.pad(x, mode='edge', pad_width=1) # pad_width -> search radius
print(padded)

[[ 5  5  0  3  3  7  7]
 [ 5  5  0  3  3  7  7]
 [ 9  9  3  5  2  4  4]
 [ 7  7  6  8  8 10 10]
 [ 1  1  6  7  7  8  8]
 [ 1  1  6  7  7  8  8]]

然后我们可以堆叠所有的邻居。 这是我想概括的操作

blocked = np.dstack([
    padded[0:-2, 0:-2], # upper left
    padded[0:-2, 1:-1], # upper center
    padded[0:-2, 2:],   # upper right
    padded[1:-1, 0:-2], # middle left...
    padded[1:-1, 1:-1],
    padded[1:-1, 2:],
    padded[2:, 0:-2],   # lower left ...
    padded[2:, 1:-1],
    padded[2:, 2:],
])

如果单元格看起来像这样，则访问邻居（对 reshape 的调用仅用于说明目的）

print(blocked[2, 2, :].reshape(3, 3))
[[3 5 2]
 [6 8 8]
 [6 7 7]]

主要问题

对于给定的搜索半径，有没有一种有效的方法来概括对np.dstack的调用？

Answer 1

这可能是一种方法 -

import numpy as np

# Parameters
R = 3  # Radius
M1,N1 = padded.shape
rowlen = N1 - R + 1
collen = M1 - R + 1

# Linear indices for the starting R x R block
idx1 = np.arange(R)[:,None]*N1 + np.arange(R)

# Offset (from the starting block indices) linear indices for all the blocks
idx2 = np.arange(collen)[:,None]*N1 + np.arange(rowlen)

# Finally, get the linear indices for all blocks
all_idx = idx1.ravel()[None,None,:] + idx2[:,:,None]

# Index into padded for the final output
out = padded.ravel()[all_idx] 

这是半径的示例运行，R = 4 -

In [259]: padded
Out[259]: 
array([[ 5,  5,  0,  3,  3,  3],
       [ 5,  5,  0,  3,  3,  3],
       [ 7,  7,  9,  3,  5,  5],
       [ 2,  2,  4,  7,  6,  6],
       [ 8,  8,  8, 10,  1,  1],
       [ 6,  6,  7,  7,  8,  8],
       [ 6,  6,  7,  7,  8,  8]])

In [260]: out
Out[260]: 
array([[[ 5,  5,  0,  3,  5,  5,  0,  3,  7,  7,  9,  3,  2,  2,  4,  7],
        [ 5,  0,  3,  3,  5,  0,  3,  3,  7,  9,  3,  5,  2,  4,  7,  6],
        [ 0,  3,  3,  3,  0,  3,  3,  3,  9,  3,  5,  5,  4,  7,  6,  6]],

       [[ 5,  5,  0,  3,  7,  7,  9,  3,  2,  2,  4,  7,  8,  8,  8, 10],
        [ 5,  0,  3,  3,  7,  9,  3,  5,  2,  4,  7,  6,  8,  8, 10,  1],
        [ 0,  3,  3,  3,  9,  3,  5,  5,  4,  7,  6,  6,  8, 10,  1,  1]],

       [[ 7,  7,  9,  3,  2,  2,  4,  7,  8,  8,  8, 10,  6,  6,  7,  7],
        [ 7,  9,  3,  5,  2,  4,  7,  6,  8,  8, 10,  1,  6,  7,  7,  8],
        [ 9,  3,  5,  5,  4,  7,  6,  6,  8, 10,  1,  1,  7,  7,  8,  8]],

       [[ 2,  2,  4,  7,  8,  8,  8, 10,  6,  6,  7,  7,  6,  6,  7,  7],
        [ 2,  4,  7,  6,  8,  8, 10,  1,  6,  7,  7,  8,  6,  7,  7,  8],
        [ 4,  7,  6,  6,  8, 10,  1,  1,  7,  7,  8,  8,  7,  7,  8,  8]]])