Ускорение кода на Python средствами самого языка

Каким бы хорошим не был Python, есть у него проблема известная все разработчикам — скорость. На эту тему было написано множество статей, в том числе и на Хабре.


Чаще всего, предлагают следующие решения:

• Использовать Psyco
• Переписать часть программы на С используя Python C Extensions
• Сменить мозгиалгоритм

Безусловно, решения верные. Но у каждого из них есть свои недостатки.
Psyco — прекрасный модуль, достигающий ускорения кода в сотни процентов, но: поддерживаются лишь 32-битный Python версий не выше 2.6, большое потребление памяти и тот факт, что в последнее время разработка psyco сбавила темпы, последнее обновление на сайте датировано 16.07.2010. Возможно, с выходом Psyco v2 ситуация изменится, но пока что, этот модуль применим не всегда.
Python C Extensions (рекомендую отличную статью rushman Пишем модуль расширения для Питона на C) — создатели Python'a сделали всем разработчикам, использующим этот язык, неоценимый подарок — Python/С API, дающий возможность относитенльно прозрачной интеграции Си-шного кода в программы на Python'e. Недостатков у этого решения только два:

1. «Порог вхождения» у C и Python/C API все же выше, чем у «голого» Python'a, что отсекает эту возможность для разработчиков, не знакомых с C
2. Одной из ключевых особенностей Python является скорость разработки. Написание части программы на Си снижает ее, пропорционально части переписанного в Си кода к всей программе

Так что, данный метод тоже подойдет не всем.
Смена алгоритма же возможна не всегда, часты ситуации, что и самый быстрый алгоритм выдает разочаровывающие результаты по скорости.

Так что же делать?

Тогда, если для вашего проекта выше перечисленные методы не подошли, что делать? Менять Python на другой язык? Нет, сдаваться нельзя. Будем оптимизировать сам код. Примеры будут взяты из программы, строящей множество Мандельброта заданного размера с заданным числом итераций.
Время работы исходной версии при параметрах 600*600 пикселей, 100 итераций составляло 3.07 сек, эту величину мы возьмем за 100%

Скажу заранее, часть оптимизаций приведет к тому, что код станет менее pythonic, да простят меня адепты python-way.

Шаг 0. Вынос основного кода программы в отдельную

Данный шаг помогает интерпретатору python лучше проводить внутренние оптимизации про запуске, да и при использовании psyco данный шаг может сильно помочь, т.к. psyco оптимизирует лишь функции, не затрагивая основное тело программы.
Если раньше рассчетная часть исходной программы выглядела так:

for Y in xrange(height):
        for X in xrange(width):
                #проверка вхождения точки (X,Y) в множество Мандельброта, itt итераций 

То, изменив её на:

def mandelbrot(height, itt, width):
    for Y in xrange(height):
        for X in xrange(width):
             #проверка вхождения точки (X,Y) в множество Мандельброта, itt итераций 
mandelbrot(height, itt, width)

мы получили время 2.4 сек, т.е. 78% от исходного.

Шаг 1. Профилирование

Стандартная библиотека Python'a, это просто клондайк полезнейших модулей. Сейчас нас интересует модуль cProfile, благодаря которому, профилирование кода становится простым и, даже, интересным занятием.
Полную документацию по этому модулю можно найти здесь, нам же хватит пары простых команд.

python -m cProfile sample.py
Ключ интерпетатора -m позволяет запускать модули как отдельные программы, если сам модуль предоставляет такую возможность.
Результатом этой команды будет получение «профиля» программы — таблицы, вида
4613944 function calls (4613943 primitive calls) in 2.818 seconds

Ordered by: internal time

ncalls tottime percall cumtime percall filename:lineno(function)
1 2.309 2.309 2.766 2.766 mand_slow.py:22(mandelbrot)
...

С её помощью, легко определить места, требующие оптимизации (строки с наибольшими значениями ncalls (кол-во вызовов функции), tottime и percall (время работы всех вызовов данной функции и каждого отдельного соответственно)).

Для удобства можно добавить ключ -s time, отсортировав вывод профилировщика по времени выполнения.

В моем случае интересной частью вывода было (время выполнение отличается от указанного выше, т.к. профилировщик добавляет свой «оверхед»):
4613944 function calls (4613943 primitive calls) in 2.818 seconds

Ordered by: internal time

ncalls tottime percall cumtime percall filename:lineno(function)
1 2.309 2.309 2.766 2.766 mand_slow.py:22(mandelbrot)
3533224 0.296 0.000 0.296 0.000 {abs}
360000 0.081 0.000 0.081 0.000 {math.atan2}
360000 0.044 0.000 0.044 0.000 {math.cos}
360000 0.036 0.000 0.036 0.000 {math.sqrt}
...
Итак, профиль получен, теперь займемся оптимизацией вплотную.

Шаг 2. Анализ профиля

Видим, что на первом месте по времени стоит наша основная функция mandelbrot, за ней идет системная функция abs, за ней несколько функций из модуля math, далее — одиночные вызовы функций, с минимальными временными затратами, они нам не интересны.

Итак, системные функции, «вылизаные» сообществом, нам врядли удастся улучшить, так что перейдем к нашему собственному коду:

Шаг 3. Математика

Сейчас, код выглядит так:

pix = img.load() #загрузим массив пикселей
def mandelbrot(height, itt, width):
    step_x = (2 - width / 1.29) / (width / 2.6) - (1 - width / 1.29) / (width / 2.6) #шаг по оси х 
    for Y in xrange(height):
        y = (Y - height / 2) / (width / 2.6) #для Y рассчет шага не так критичен как для Х, его отсутствие положительно повлияет на точность
        x = - (width / 1.29) / (width / 2.6)
        for X in xrange(width):
            x += step_x
            z = complex(x, y)

            phi = math.atan2(y, x - 0.25)
            p = math.sqrt((x - 0.25) ** 2 + y ** 2)
            pc = 0.5 - 0.5 * math.cos(phi)

            if p <= pc: #если лежит в области кардиоиды - отмечаем
                pix[X, Y] = (255, 255, 255)
                continue

            Z_i = 0j
            for i in xrange(itt): #проверка на выход за "границы бесконечности"
                Z_i = Z_i ** 2 + z
                if abs(Z_i) > 2:
                    color = (i * 255) // itt
                    pix[X, Y] = (color, color, color)
                    break
            else:
                pix[X, Y] = (255, 255, 255)
        print("\r%d/%d" % (Y, height)),

Заметим, что оператор возведения в степень ** — довольно «общий», нам же необходимо лишь возведение во вторую степень, т.е. все конструкции вида x**2 можно заменить на х*х, выиграв таким образом еще немного времени. Посмотрим на время:
1.9 сек, или 62% изначального времени, достигнуто простой заменой двух строк:

p = math.sqrt((x - 0.25) ** 2 + y ** 2)
...
Z_i = Z_i **2 + z

на:

p = math.sqrt((x - 0.25) * (x - 0.25) + y * y)
...
Z_i = Z_i * Z_i + z

Шажки 5, 6 и 7. Маленькие, но важные

Прописная истина, о которой знают все программисты на Python — работа с глобальными переменными медленней работы с локальными. Но часто забывается факт, что это верно не только для переменных но и вообще для всех объектов. В коде функции идут вызовы нескольких функций из модуля math. Так почему бы не импортировать их в самой функции? Сделано:

def mandelbrot(height, itt, width):
    from math import atan2, cos, sqrt
    pix = img.load() #загрузим массив пикселей

Еще 0.1сек отвоевано.
Вспомним, что abs(x) вернет число типа float. Так что и сравнивать его стоит с float а не int:

if abs(Z_i) > 2:  ------> if abs(Z_i) > 2.0: 

Еще 0.15сек. 53% от начального времени.

И, наконец, грязный хак.
В конкретно этой задаче, можно понять, что нижняя половина изображения, равна верхней, т.о. число вычислений можно сократить вдвое, получив в итоге 0.84сек или 27% от исходного времени.

Заключение

Профилируйте. Используйте timeit. Оптимизируйте. Python — мощный язык, и программы на нем будут работать со скоростью, пропорциональной вашему желанию разобраться и все отполировать:)
Цель данной статьи, показать, что за счет мелких и незначительных изменения, таких как замен ** на *, можно заставить зеленого змея ползать до двух раз быстрее, без применения тяжелой артиллерии в виде Си, или шаманств psyco.
Также, можно совместить разные средства, такие как вышеуказанные оптимизации и модуль psyco, хуже не станет:)

Спасибо всем кто дочитал до конца, буду рад выслушать ваши мнения и замечания в комментариях!

UPD Полезную ссылку в коментариях привел funca.