Python向量化嵌套的for循环

Question

我希望您能协助找到并理解一种Pythonic的方式来优化以下嵌套for循环中的数组操作：

def _func(a, b, radius):
    "Return 0 if a>b, otherwise return 1"
    if distance.euclidean(a, b) < radius:
        return 1
    else:
        return 0

def _make_mask(volume, roi, radius):
    mask = numpy.zeros(volume.shape)
    for x in range(volume.shape[0]):
        for y in range(volume.shape[1]):
            for z in range(volume.shape[2]):
                mask[x, y, z] = _func((x, y, z), roi, radius)
    return mask

当 volume.shape（182, 218, 200）和 roi.shape（3）都是 ndarray 类型时；而 radius 是一个 int

- Kambiz

---

这些答案中有帮助您的吗？请阅读相关页面了解：接受答案是如何工作的？ - rayryeng

请原谅我发表这篇旧帖，但是 A：你应该接受 @Divakar 的帖子。它是使用 numpy 进行向量化的精彩演示，B：你应该查看 scipy.spatial 中的 KD 树和球点算法。当数据稀疏或不在常规网格上时，这是一种可推广的方法来解决你的特定问题。虽然它不是这个确切问题的最佳方法，但了解它是非常好的（我最近自己使用了它）。 - Aaron

@Divakar，你的解释非常有帮助，谢谢。我最初已经点赞了，但最近才意识到勾选标记的目的。这件事已经完成了。 - Kambiz

Answer 1

方法 #1

这是一种矢量化的方法 -

m,n,r = volume.shape
x,y,z = np.mgrid[0:m,0:n,0:r]
X = x - roi[0]
Y = y - roi[1]
Z = z - roi[2]
mask = X**2 + Y**2 + Z**2 < radius**2

可能的改进：我们可以使用numexpr模块来加速最后一步 -

import numexpr as ne

mask = ne.evaluate('X**2 + Y**2 + Z**2 < radius**2')

第二种方法

我们还可以逐步构建与形状参数相对应的三个范围，并在执行减法时对roi的三个元素进行即时操作，而不像之前使用np.mgrid一样实际创建网格。为了提高效率，这可以通过使用broadcasting来实现。实现代码如下:

m,n,r = volume.shape
vals = ((np.arange(m)-roi[0])**2)[:,None,None] + \
       ((np.arange(n)-roi[1])**2)[:,None] + ((np.arange(r)-roi[2])**2)
mask = vals < radius**2

简化版本：感谢@Bi Rico在这里提出的改进建议，我们可以使用np.ogrid以更加简明的方式执行这些操作，如下所示 -

m,n,r = volume.shape    
x,y,z = np.ogrid[0:m,0:n,0:r]-roi
mask = (x**2+y**2+z**2) < radius**2

---

运行时测试

函数定义 -

def vectorized_app1(volume, roi, radius):
    m,n,r = volume.shape
    x,y,z = np.mgrid[0:m,0:n,0:r]
    X = x - roi[0]
    Y = y - roi[1]
    Z = z - roi[2]
    return X**2 + Y**2 + Z**2 < radius**2

def vectorized_app1_improved(volume, roi, radius):
    m,n,r = volume.shape
    x,y,z = np.mgrid[0:m,0:n,0:r]
    X = x - roi[0]
    Y = y - roi[1]
    Z = z - roi[2]
    return ne.evaluate('X**2 + Y**2 + Z**2 < radius**2')

def vectorized_app2(volume, roi, radius):
    m,n,r = volume.shape
    vals = ((np.arange(m)-roi[0])**2)[:,None,None] + \
           ((np.arange(n)-roi[1])**2)[:,None] + ((np.arange(r)-roi[2])**2)
    return vals < radius**2

def vectorized_app2_simplified(volume, roi, radius):
    m,n,r = volume.shape    
    x,y,z = np.ogrid[0:m,0:n,0:r]-roi
    return (x**2+y**2+z**2) < radius**2

时间 -

In [106]: # Setup input arrays  
     ...: volume = np.random.rand(90,110,100) # Half of original input sizes 
     ...: roi = np.random.rand(3)
     ...: radius = 3.4
     ...: 

In [107]: %timeit _make_mask(volume, roi, radius)
1 loops, best of 3: 41.4 s per loop

In [108]: %timeit vectorized_app1(volume, roi, radius)
10 loops, best of 3: 62.3 ms per loop

In [109]: %timeit vectorized_app1_improved(volume, roi, radius)
10 loops, best of 3: 47 ms per loop

In [110]: %timeit vectorized_app2(volume, roi, radius)
100 loops, best of 3: 4.26 ms per loop

In [139]: %timeit vectorized_app2_simplified(volume, roi, radius)
100 loops, best of 3: 4.36 ms per loop

所以，就像往常一样，broadcasting 展示了它的魔力，使得代码速度疯狂提升了近10,000倍，比使用即时广播操作创建网格要好10倍以上！

- Divakar

---

方法二与方法一非常相似，只是使用 np.ogrid 替换了 np.mgrid。 - Bi Rico

我们能否使用 ogrid 而不是 mgrid 来获取 'app1' 的时间呢 :)。 - Bi Rico

@BiRico 为什么要选择“取其一部分”，我们可以得到所有的东西 :) 非常感谢您的改进，现在看起来更加简洁！ - Divakar

Answer 2

假设你首先创建了一个 xyzy 数组：

import itertools

xyz = [np.array(p) for p in itertools.product(range(volume.shape[0]), range(volume.shape[1]), range(volume.shape[2]))]

现在，使用 numpy.linalg.norm，

np.linalg.norm(xyz - roi, axis=1) < radius

检查每个元组从 roi 的距离是否小于半径。

最后，只需将结果 reshape 成所需的维度。

- Ami Tavory