Векторизация Python, вложенная для циклов
Я буду благодарна за помощь в поиске и понимании типичный для Python способ оптимизировать следующие манипуляции массив вложенных циклов 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
2 ответов
подход #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
мы также можем постепенно построить три диапазона, соответствующие параметрам формы и выполнить вычитание против трех элементов 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,000x
ускорение над исходным кодом и более чем 10x
лучше, чем создание сеток с помощью на лету транслируемых операций!
скажем, вы сначала построите 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
результат к размерам вам.