时间戳

首页 > 解决方案 > 在 numpy ndarray 中进行高效的邻域搜索，而不是嵌套的条件 for 循环

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

问题描述

尽管有很多问题实例：“嵌套 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) 节点： 原始矩阵和掩码矩阵的样本切片

标签： pythonperformancenumpynested-loopsneighbours

解决方案

由于您的循环仅使用直接 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 实现快 362 倍。

请注意，compute_imtemp由于编译，第一次调用会很慢。克服这个问题的一种方法是调用compute_imtemp一个空的 Numpy 数组。另一种方法是使用 Numba API 手动编译函数并将类型显式提供给 Numba。

推荐阅读