Python: векторизация циклов с условиями дает ускорение примерно на 20—50%

Одним из приемов для ускорения работы циклов является векторизация вычислений, т. е. использование функций, которые поддерживают операции над векторами. Лучший способ ускорить цикл — отказаться от него. В статье сделан акцент на векторизации условных циклов и продемонстрирован пример с результатами измерения производительности.

Математические преобразования, основанные на некоторых предопределенных условиях, хорошо распространены в научных задачах. И получается, что можно легко векторизовать простые блоки условных циклов, сначала превратив их в функции, а затем использовать метод numpy.vectorize().

Следующий пример удобно запускать в Jypiter Notebook.

Пример

import numpy as np
from math import sin as sn
import matplotlib.pyplot as plt
import time

# Количество точек тестирования
N_point  = 1000

# Пользовательская функция с условными циклами
def myfunc(x,y):
    if (x>0.5*y and y<0.3):
        return (sn(x-y))
    elif (x<0.5*y):
        return 0
    elif (x>0.2*y):
        return (2*sn(x+2*y))
    else:
        return (sn(y+x))

# Список элементов, сфомированных по закону нормального распределения
lst_x = np.random.randn(N_point)
lst_y = np.random.randn(N_point)
lst_result = []

# Визуализация в виде графиков
plt.hist(lst_x,bins=20)
plt.show()
plt.hist(lst_y,bins=20)
plt.show()

# Сначала чистый for-loop
t1=time.time()
for i in range(len(lst_x)):
    x = lst_x[i]
    y= lst_y[i]
    if (x>0.5*y and y<0.3):
        lst_result.append(sn(x-y))
    elif (x<0.5*y):
        lst_result.append(0)
    elif (x>0.2*y):
        lst_result.append(2*sn(x+2*y))
    else:
        lst_result.append(sn(y+x))
t2=time.time()
print("\nВремя, затраченное обычным for-loop\n{}\n{} us".format('-'*47,1000000*(t2-t1)))

# Списковое выражение
print("\nВремя, затраченное списковым выражением и zip()\n"+'-'*47)
%timeit lst_result = [myfunc(x,y) for x,y in zip(lst_x,lst_y)]

# Функция map()
print("\nВремя, затраченное функцией map()\n"+'-'*47)
%timeit list(map(myfunc,lst_x,lst_y))

# Метод Numpy.vectorize()
print("\nВремя, затраченное numpy.vectorize()\n"+'-'*47)
vectfunc = np.vectorize(myfunc,otypes=[np.float],cache=False)
%timeit list(vectfunc(lst_x,lst_y))

Результаты

Время, затраченное обычным for-loop
-----------------------------------------------
2000.0934600830078 us

Время, затраченное списковым выражением и zip()
-----------------------------------------------
1000 loops, best of 3: 810 µs per loop

Время, затраченное функцией map()
-----------------------------------------------
1000 loops, best of 3: 726 µs per loop

Время, затраченное numpy.vectorize()
-----------------------------------------------
1000 loops, best of 3: 516 µs per loop

Обратите внимание на то, что в коде выше постоянно используется вызов time.time(). Таким образом, эффективно запускается не менее 1000 циклов одного и того же выражения и усредняется время выполнения, чтобы избежать влияния случайных факторов.

Мы видим доказательства того, что для задачи, основанной на серии условных проверок, векторизация с использованием numpy дает ускорение примерно на 20—50% по сравнению с обычными Python-методами. Может быть, выигрыш в 20—50% не впечатлит вас, но каждый бит экономии времени складывается в процессе передачи данных и возвращается в конечном итоге! Если задача требует, чтобы подобное преобразование произошло миллион раз, может возникнуть разница между 2 днями и 8 часами.

Выводы

Если у вас есть длинный список данных и нужно выполнить какое-либо математическое преобразование над ними, обязательно рассмотрите возможность превращения этих структур данных Python (списки, кортежи или словари) в объекты numpy.ndarray и используйте возможности встроенной векторизации.

Numpy предоставляет C-API для еще более быстрого выполнения кода, но это усложняет программирование на Python.

Кстати, по теме статьи есть отличная онлайн-книга с открытым исходным кодом от французского исследователя нейробиологии: https://www.labri.fr/perso/nrougier/from-python-to-numpy/#id7.


По материалам Tirthajyoti Sarkar.