numpy-ndarray代替嵌套条件for循环的有效邻域搜索

Question

尽管有很多这样的问题：“嵌套for循环的numpy替代方案是什么”，但我无法找到适合我的案例的答案。下面是：

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

示例输入数据：

import numpy as np

im = np.array([[[ 60, 54, 47, 52, 57, 53, 46, 48]
, [ 60, 57, 53, 53, 54, 53, 50, 55]
, [ 60, 63, 56, 58, 59, 57, 50, 50]
, [ 70, 70, 64, 69, 74, 72, 64, 47]
, [ 73, 76, 77, 80, 82, 76, 58, 37]
, [ 85, 85, 86, 86, 78, 62, 38, 20]
, [ 94, 94, 92, 78, 54, 33, 16, 255]
, [ 94, 90, 72, 51, 32, 19, 255, 255]
, [ 65, 53, 29, 18, 255, 255, 255, 255]
, [ 29, 22, 255, 255, 255, 255, 255,  0]]

, [[ 66, 67, 70, 69, 75, 73, 72, 63]
, [ 68, 70, 73, 74, 78, 80, 74, 53]
, [ 75, 87, 87, 83, 89, 86, 61, 33]
, [ 81, 89, 88, 98, 99, 77, 41, 18]
, [ 84, 94, 100, 100, 82, 49, 21, 255]
, [ 99, 101, 92, 75, 48, 25, 255, 255]
, [ 93, 77, 52, 32, 255, 255, 255, 255]
, [ 52, 40, 25, 255, 255, 255, 255, 255]
, [ 23, 16, 255, 255, 255, 255, 255,  0]
, [255, 255, 255, 255, 255, 255,  0,  0]]

, [[ 81, 83, 92, 101, 101, 83, 49, 19]
, [ 86, 96, 103, 103, 95, 64, 28, 255]
, [ 94, 103, 107, 98, 79, 41, 255, 255]
, [101, 103, 98, 79, 51, 28, 255, 255]
, [102, 97, 76, 49, 27, 255, 255, 255]
, [ 79, 62, 35, 21, 255, 255, 255, 255]
, [ 33, 23, 15, 255, 255, 255, 255, 255]
, [ 16, 255, 255, 255, 255, 255, 255,  0]
, [255, 255, 255, 255, 255, 255,  0,  0]
, [255, 255, 255, 255, 255,  0,  0,  0]]

, [[106, 107, 109, 94, 58, 26, 15, 255]
, [110, 104, 90, 66, 37, 19, 255, 255]
, [106, 89, 61, 35, 22, 255, 255, 255]
, [ 76, 56, 34, 19, 255, 255, 255, 255]
, [ 40, 27, 18, 255, 255, 255, 255, 255]
, [ 17, 255, 255, 255, 255, 255, 255, 255]
, [255, 255, 255, 255, 255, 255, 255,  0]
, [255, 255, 255, 255, 255, 255,  0,  0]
, [255, 255, 255, 255, 255,  0,  0,  0]
, [255, 255, 255,  0,  0,  0,  0,  0]]

, [[ 68, 51, 33, 19, 255, 255, 255, 255]
, [ 45, 34, 20, 255, 255, 255, 255, 255]
, [ 28, 18, 255, 255, 255, 255, 255, 255]
, [ 17, 255, 255, 255, 255, 255, 255, 255]
, [255, 255, 255, 255, 255, 255, 255, 255]
, [255, 255, 255, 255, 255, 255, 255,  0]
, [255, 255, 255, 255, 255, 255,  0,  0]
, [255, 255, 255, 255, 255,  0,  0,  0]
, [255, 255, 255,  0,  0,  0,  0,  0]
, [255,  0,  0,  0,  0,  0,  0,  0]]

, [[255, 255, 255, 255, 255, 255, 255, 255]
, [255, 255, 255, 255, 255, 255, 255, 255]
, [255, 255, 255, 255, 255, 255, 255, 255]
, [255, 255, 255, 255, 255, 255, 255,  0]
, [255, 255, 255, 255, 255, 255,  0,  0]
, [255, 255, 255, 255, 255,  0,  0,  0]
, [255, 255, 255, 255,  0,  0,  0,  0]
, [255, 255, 255,  0,  0,  0,  0,  0]
, [255,  0,  0,  0,  0,  0,  0,  0]
, [  0,  0,  0,  0,  0,  0,  0,  0]]

, [[255, 255, 255, 255, 255, 255, 255,  0]
, [255, 255, 255, 255, 255, 255, 255,  0]
, [255, 255, 255, 255, 255, 255,  0,  0]
, [255, 255, 255, 255, 255,  0,  0,  0]
, [255, 255, 255, 255,  0,  0,  0,  0]
, [255, 255, 255,  0,  0,  0,  0,  0]
, [255, 255,  0,  0,  0,  0,  0,  0]
, [  0,  0,  0,  0,  0,  0,  0,  0]
, [  0,  0,  0,  0,  0,  0,  0,  0]
, [  0,  0,  0,  0,  0,  0,  0,  0]]

, [[255, 255, 255, 255, 255, 255,  0,  0]
, [255, 255, 255, 255, 255,  0,  0,  0]
, [255, 255, 255, 255,  0,  0,  0,  0]
, [255, 255, 255,  0,  0,  0,  0,  0]
, [255, 255,  0,  0,  0,  0,  0,  0]
, [255,  0,  0,  0,  0,  0,  0,  0]
, [  0,  0,  0,  0,  0,  0,  0,  0]
, [  0,  0,  0,  0,  0,  0,  0,  0]
, [  0,  0,  0,  0,  0,  0,  0,  0]
, [  0,  0,  0,  0,  0,  0,  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,j,k]=-2
                elif imtemp[i+1,j,k] == 0 or imtemp[i-1,j,k] == 0 or imtemp[i,j+1,k] == 0 or imtemp[i,j-1,k] == 0 or imtemp[i,j,k+1] == 0 or imtemp[i,j,k-1] == 0:
                    imtemp[i,j,k]=-2
                    
LA = np.argwhere(imtemp==-2)        

上述示例代码产生的LA为：

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

以及Z方向的切片（XY平面实例），显示不同的未接触、遮罩（-1）和目标（-2）节点：

Answer 1

问题陈述

在一个大数组中有一个“实体”形状。你从上面刻出一个球。你的目标是找到球内固体表面的指数。曲面定义为具有6点连接的实体外部相邻的任何点。阵列的边也被视为曲面

更快的环路解决方案

您已经计算了表示实体和球的交点的遮罩。您可以更优雅地计算掩码，并将其转换为索引。我建议保持维度的顺序不变，而不是在不同的符号之间切换。例如，RN的顺序会受到影响，并且您会面临与轴限制不匹配的风险

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

im = ...

cutout = ((np.indices(im.shape) - RN.reshape(-1, 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), 0, 4) - RN)**2).sum(axis=-1) <= rad**2

计算索引的好处在于，循环不再需要太大。通过使用argwhere，基本上去掉了外部的三个循环，只剩下if语句进行循环。您还可以将连接检查矢量化。这有一个很好的副作用，你可以为每个像素定义任意的连接

limit = np.array(im.shape) - 1  # Edge of `im`
connectivity = np.array([[ 1,  0,  0],  # Add rows to determine connectivity
                         [-1,  0,  0],
                         [ 0,  1,  0],
                         [ 0, -1,  0],
                         [ 0,  0,  1],
                         [ 0,  0, -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。当元素通过您的标准时，您可以将元素设置为False，而不是将它们设置为-2，如果它们没有通过标准，则可以将元素设置为False

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

检查ind.all() and (ind < limit).all() and im[tuple((ind + connectivity).T)].all()是对or条件执行操作的快捷方式，但相反（测试非曲面而非曲面）

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

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

index_mask = np.zeros(len(indices), dtype=bool)

for n, 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, :]

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

完全矢量化解决方案

上面的循环可以使用广播的魔力完全矢量化。我们可以将connectivity批量添加到connectivity中并批量过滤结果，而不是在indices中的每一行上循环。诀窍是添加足够的维度，以便将所有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, 4, 3])
rad = 3

im = ...

# Find subset of interest
cutout = ((np.indices(im.shape) - RN.reshape(-1, 1, 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, None, :] + connectivity[None, ...]

# 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, :]

Answer 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,j,k]=-2
                    elif imtemp[i+1,j,k] == 0 or imtemp[i-1,j,k] == 0 or imtemp[i,j+1,k] == 0 or imtemp[i,j-1,k] == 0 or imtemp[i,j,k+1] == 0 or imtemp[i,j,k-1] == 0:
                        imtemp[i,j,k]=-2

[...]
imtemp = im.copy()
imtemp[mask] = -1
compute_imtemp(imtemp, X, Y, 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实现的速度比eem快362倍

请注意，由于编译的原因，对compute_imtemp的第一次调用将很慢。克服此问题的一种方法是在空Numpy数组上调用compute_imtemp。另一种方法是使用NumbaAPI手动编译函数，并显式地向Numba提供类型