在numpy ndarray中进行有效的邻域搜索，而不是嵌套条件for循环

如何解决在numpy ndarray中进行有效的邻域搜索，而不是嵌套条件for循环

尽管问题的实例很多：“嵌套循环的numpy替代方案是什么”，但我找不到适合我的情况的答案。在这里：

我有一个3D numpy数组，背景为“ 0”，其他整数为前景。我想找到并存储落入预定义蒙版（定义了与参考节点的给定距离的球体）内的前景体素。我已经使用嵌套的“ for”循环和一系列“ if”条件成功完成了任务，如下所示。我正在寻找一种更高效，更紧凑的替代方案，以避免这种邻域搜索算法的循环和冗长的条件。

样本输入数据：

import numpy as np

im = np.array([[[ 60,54,47,52,57,53,46,48],[ 60,50,55],63,56,58,59,50],[ 70,70,64,69,74,72,47],[ 73,76,77,80,82,37],[ 85,85,86,78,62,38,20],[ 94,94,92,33,16,255],90,51,32,19,255,[ 65,29,18,[ 29,22,0]],[[ 66,67,75,73,63],[ 68,53],[ 75,87,83,89,61,33],[ 81,88,98,99,41,18],[ 84,100,49,21,[ 99,101,48,25,[ 93,[ 52,40,[ 23,0],[255,[[ 81,19],[ 86,96,103,95,28,107,79,[101,[102,97,27,[ 79,35,[ 33,23,15,[ 16,[[106,109,26,[110,104,66,37,[106,[ 76,34,[ 40,[ 17,[[ 68,[ 45,20,[ 28,[[255,[  0,0]]])

已实现的方法：

[Z,Y,X]=im.shape
RN = np.array([3,4,4])     
################Loading Area search
rad = 3
a,b,c = RN
x,y,z = np.ogrid[-c:Z-c,-b:Y-b,-a:X-a]
neighborMask = x*x + y*y + z*z<= rad*rad
noNodeMask = im > 0
mask = np.logical_and(neighborMask,noNodeMask)

imtemp = im.copy()
imtemp[mask] = -1

for i in range (X):
    for j in range (Y):
        for k in range (Z):
            if imtemp[i,j,k]==-1:
                if i in (0,X-1) or j in (0,Y-1) or k in (0,Z-1): 
                    imtemp[i,k]=-2
                elif imtemp[i+1,k] == 0 or imtemp[i-1,k] == 0 or imtemp[i,j+1,j-1,k+1] == 0 or imtemp[i,k-1] == 0:
                    imtemp[i,k]=-2
                    
LA = np.argwhere(imtemp==-2)        

上面的示例代码生成的LA为：

In [90]:LA
Out[90]: 
array([[4,[4,6],5,5],6,4],7,3],[5,3,2],[6,2,1],2]])

以及在Z方向上的一个切片（XY平面实例），该切片显示了不同的未触摸，蒙版（-1）和目标（-2）节点：

解决方法

由于循环仅使用直接的Numpy索引，因此可以使用 Numba的@njit 以更有效的方式执行此操作。

@njit
def compute_imtemp(imtemp,X,Y,Z):
    for i in range (Z):
        for j in range (Y-1):
            for k in range (X-1):
                if imtemp[i,j,k]==-1:
                    if i==(Z-1): 
                        imtemp[i,k]=-2
                    elif imtemp[i+1,k] == 0 or imtemp[i-1,k] == 0 or imtemp[i,j+1,j-1,k+1] == 0 or imtemp[i,k-1] == 0:
                        imtemp[i,k]=-2

[...]
imtemp = im.copy()
imtemp[mask] = -1
compute_imtemp(imtemp,Z)
LA = np.argwhere(imtemp==-2)

这是我的机器上的性能结果：

281 µs ± 1.43 µs per loop (mean ± std. dev. of 7 runs,100 loops each)
776 ns ± 16.4 ns per loop (mean ± std. dev. of 7 runs,100 loops each)

Numba的实现速度快了 362倍。

请注意，由于编译，第一次调用compute_imtemp会很慢。解决此问题的一种方法是在空的Numpy数组上调用compute_imtemp。另一种方法是使用Numba API手动编译该函数，并将类型明确提供给Numba。

,

问题陈述

您的阵列呈“实心”形状。你从中挖出一个球。您的目标是找到球内实体表面的索引。曲面定义为与实体外部相邻的，具有6点连通性的任何点。数组的边缘也被视为曲面。

更快的循环解决方案

您已经计算了代表实体与球的交点的蒙版。您可以更优雅地计算蒙版，然后将其转换为索引。我建议保持尺寸顺序不变，而不要在不同的符号之间切换。例如，RN的顺序会受到影响，并且冒着轴限制不匹配的风险。

RN = np.array([4,4,3])
rad = 3

im = ...

cutout = ((np.indices(im.shape) - RN.reshape(-1,1,1))**2).sum(axis=0) <= rad**2
solid = im > 0
mask = solid & cutout
indices = np.argwhere(mask)

您也可以在不重塑RN的情况下获取抠图

cutout = ((np.rollaxis(np.indices(im.shape,sparse=False),4) - RN)**2).sum(axis=-1) <= rad**2

关于计算索引的好处是您的循环不再需要很大。通过使用argwhere，您基本上去除了外部的三个循环，仅留下if语句进行循环。您还可以矢量化连接检查。这样做的好处是，您可以为每个像素定义任意连接。

limit = np.array(im.shape) - 1  # Edge of `im`
connectivity = np.array([[ 1,0],# Add rows to determine connectivity
                         [-1,[ 0,-1,1],-1]],dtype=indices.dtype)
index_mask = np.ones(len(indices),dtype=bool)

for n,ind in enumerate(indices):
    if ind.all() and (ind < limit).all() and im[tuple((ind + connectivity).T)].all():
        index_mask[n] = False

LA = indices[index_mask,:]

请注意，imtemp根本没有意义。即使在原始循环中，您也可以直接操作mask。如果元素不符合条件，则可以将元素设置为-2，而不是将元素设置为False。

我在这里做类似的事情。我们检查实际选择的每个索引，并确定其中是否有任何一个在实体内。这些指标从遮罩中消除。然后根据掩码更新索引列表。

支票ind.all() and (ind < limit).all() and im[tuple((ind + connectivity).T)].all()是您在or条件下执行操作的捷径，但相反（检查非表面而不是表面）。

• ind.all()检查是否所有索引都不为零：即不在顶部/顶部/左侧。
• (ind < limit).all()检查是否所有索引均不等于相应的图像大小减去一。
• im[tuple((ind + connectivity).T)].all()检查是否所有连接的像素都不为零。 (ind + connectivity).T是我们连接到的六个点的(3,6)数组（当前(6,3) connectivity数组在每个轴上定义为+/- 1）。当您将其变成元组时，它就变成了精美的索引，就好像您做了im[x + connectivity[:,y + connectivity[:,z + connectivity[:,2]]之类的事情一样。索引中的逗号只是使其成为一个元组。我展示的方式更适合任意数量的尺寸。

通过所有三个测试的像素位于实体内部，并被移除。当然，您可以编写循环以另一种方式进行检查，但是随后您必须更改掩码：

index_mask = np.zeros(len(indices),ind in enumerate(indices):
    if (ind == 0).any() or (ind == limit).any() or (im[tuple((ind + connectivity).T)] == 0).any():
        index_mask[n] = True

LA = indices[index_mask,:]

手动循环无论如何都不理想。但是，它向您展示了如何缩短循环（可能缩短几个数量级），以及如何使用向量化和广播定义任意连接性，而不会因对其进行硬编码而陷入困境。

完全矢量化解决方案

上面的循环可以使用广播的魔力完全矢量化。无需遍历indices中的每一行，我们可以向其中批量添加connectivity并批量过滤结果。技巧是添加足够的尺寸，以便将所有connectivity添加到indices的每个元素。

您仍将要忽略边缘的像素：

edges = (indices == 0).any(axis=-1) | (indices == limit).any(axis=-1)
conn_index = indices[~edges,None,:] + connectivity[None,...]

index_mask = np.empty(len(indices),dtype=bool)
index_mask[edges] = True
index_mask[~edges] = (im[tuple(conn_index.T)] == 0).any(axis=0)

LA = indices[index_mask,:]

我希望使用numba编译的正确编写的循环将比该解决方案快得多，因为它将避免流水线操作的大量开销。不需要大的临时缓冲区或特殊处理。

TL; DR

# Parameters
RN = np.array([4,3])
rad = 3

im = ...

# Find subset of interest
cutout = ((np.indices(im.shape) - RN.reshape(-1,1))**2).sum(axis=0) <= rad**2
solid = im > 0

# Convert mask to indices
indices = np.argwhere(solid & cutout)

# Find image edges among indices
edges = (indices == 0).any(axis=-1) | (indices == limit).any(axis=-1)

# Connectivity elements for non-edge pixels
conn_index = indices[~edges,...]

# Mask the valid surface pixels
index_mask = np.empty(len(indices),dtype=bool)
index_mask[edges] = True
index_mask[~edges] = (im[tuple(conn_index.T)] == 0).any(axis=0)

# Final result
LA = indices[index_mask,:]