Python向量化嵌套循环

如果您能帮助我找到并理解一种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

定义函数（a、b、半径）： 如果a>b，则返回0，否则返回1 如果距离欧几里德（a，b）<半径： 返回1 其他： 返回0 def制作遮罩（体积、roi、半径）： 掩码=numpy.zero（体积.形状） 对于范围内的x（体积形状[0]）： 对于范围内的y（体积形状[1]）： 对于范围内的z（体积形状[2]）： 遮罩[x，y，z]=_func（（x，y，z），roi，半径） 返回掩码

---

其中

volume.shape

（182218200）和

roi.shape

（3，）都是

ndarray

类型；而

radius

是一个

int

假设您首先构建了一个

xyzy

数组：

import itertools

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

现在用,

np.linalg.norm（xyz-roi，轴=1）

检查每个元组距

roi

的距离是否小于半径

最后，只需将结果重塑为所需的尺寸。

方法#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

方法#2

我们还可以逐步建立与形状参数对应的三个范围，并对

roi

的三个元素执行动态减法，而无需像前面使用

np.mgrid

实际创建网格。这将得益于为提高效率目的而使用。实现如下所示-

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

---

运行时测试

函数定义-

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

---

因此，像往常一样，

broadcasting

展示了它的魔力，它比原始代码快了几近10000x，比使用即时广播操作创建网格要好得多

方法#2很像方法1，用np.ogrid替换np.mgrid。我们能用

ogrid

而不是

mgrid

：）获得“app1”的时间吗？@BiRico为什么，当我们可以得到所有东西时：）非常感谢那里的改进，现在看起来更干净了！这些答案有帮助吗？相关页面如下：请原谅我的帖子，但是A：你应该接受@Divakar的帖子。。这是一个很好的用numpy进行矢量化的演示，B：你应该看看KD树和来自的球点算法。当数据稀疏或不在规则网格上时，它是针对特定问题的通用方法。虽然这不是解决这个问题的最佳方法，但了解它（我最近也使用过它）是一件非常好的事情@divaker您的解释非常有用，谢谢。起初我投了更高的票，但最近才意识到复选标记的目的。完成了。

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

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