pythonfor循环嵌套_Python矢量化嵌套for循环

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

方法#2

我们也可以逐渐建立与形状参数相对应的三个范围，并且可以在飞行中对这三个元素进行实时减法，而不用像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

所以，像原来的代码一样，广播显示出一个疯狂的将近10,000x的速度的魔力，比通过使用即时广播操作创造网格要好10多倍！