python

Пишем код C на Cython

  • суббота, 8 ноября 2014 г. в 02:10:33
http://habrahabr.ru/company/mailru/blog/242533/

Последние два года я решаю все задачи исключительно на Cython. Это вовсе не значит, что я пишу на Питоне, а потом «Ситонизирую» это с использованием различных деклараций типов, нет, я просто пишу на Cython. Я использую «сырые» структуры и массивы C (а иногда и векторы C++) и маленькую обёртку вокруг malloc/free, которую я написал сам. Код работает практически так же быстро, как C/C++, потому что это и есть код на C/C++, украшенный синтаксическим сахаром. Это код на C/C++ с функционалом Python именно там, где мне это нужно и где я этого хочу.

Фактически это противоположный вариант стандартного применения языков, схожих с Python: вы пишете всё приложение на Питоне, оптимизируете важные места на C и… Профит! Скорость C, удобство Питона, овцы целы, волки сыты.

В теории это всегда выглядит лучше, чем на практике. На практике ваши структуры данных оказывают огромное влияние на эффективность вашего кода и трудоёмкость его написания. Работа с массивами — это всегда боль, зато они быстрые. Списки чрезвычайно удобны, но очень медленные. Циклы и вызовы функций в Питоне всегда медленные, поэтому та часть приложения, которую вы пишете на C имеет тенденцию расти и расти до тех пор, пока практически всё ваше приложение не будет написано на C.

Недавно был опубликован пост про написание расширений C для Python. Автор написал реализацию алгоритма на чистом Питоне и на C, используя Numpy C API. Я решил, что это хорошая возможность продемонстрировать различия, и, для сравнения, написал свой вариант на Cython:

import random
from cymem.cymem cimport Pool
 
from libc.math cimport sqrt
 
cimport cython
 
cdef struct Point:
    double x
    double y
 
cdef class World:
    cdef Pool mem
    cdef int N
    cdef double* m
    cdef Point* r
    cdef Point* v
    cdef Point* F
    cdef readonly double dt
    def __init__(self, N, threads=1, m_min=1, m_max=30.0, r_max=50.0, v_max=4.0, dt=1e-3):
        self.mem = Pool()
        self.N = N
        self.m = <double*>self.mem.alloc(N, sizeof(double))
        self.r = <Point*>self.mem.alloc(N, sizeof(Point))
        self.v = <Point*>self.mem.alloc(N, sizeof(Point))
        self.F = <Point*>self.mem.alloc(N, sizeof(Point))
        for i in range(N):
            self.m[i] = random.uniform(m_min, m_max)
            self.r[i].x = random.uniform(-r_max, r_max)
            self.r[i].y = random.uniform(-r_max, r_max)
            self.v[i].x = random.uniform(-v_max, v_max)
            self.v[i].y = random.uniform(-v_max, v_max)
            self.F[i].x = 0
            self.F[i].y = 0
        self.dt = dt
 
 
@cython.cdivision(True)
def compute_F(World w):
    """Compute the force on each body in the world, w."""
    cdef int i, j
    cdef double s3, tmp
    cdef Point s
    cdef Point F
    for i in range(w.N):
        # Set all forces to zero. 
        w.F[i].x = 0
        w.F[i].y = 0
        for j in range(i+1, w.N):
            s.x = w.r[j].x - w.r[i].x
            s.y = w.r[j].y - w.r[i].y
 
            s3 = sqrt(s.x * s.x + s.y * s.y)
            s3 *= s3 * s3;
 
            tmp = w.m[i] * w.m[j] / s3
            F.x = tmp * s.x
            F.y = tmp * s.y
 
            w.F[i].x += F.x
            w.F[i].y += F.y
 
            w.F[j].x -= F.x
            w.F[j].y -= F.y
 
 
@cython.cdivision(True)
def evolve(World w, int steps):
    """Evolve the world, w, through the given number of steps."""
    cdef int _, i
    for _ in range(steps):
        compute_F(w)
        for i in range(w.N):
            w.v[i].x += w.F[i].x * w.dt / w.m[i]
            w.v[i].y += w.F[i].y * w.dt / w.m[i]
            w.r[i].x += w.v[i].x * w.dt
            w.r[i].y += w.v[i].y * w.dt

Эта версия на Cython была написана за 30 минут, и она такая же быстрая, как код на C. Собственно, почему бы и нет, ведь это и есть код на C, просто написанный с применением синтаксического сахара. И вам даже не нужно думать о сложном и враждебном C API и изучать его, вы просто… просто пишете код C или C++. Обе версии, C и Cython, примерно в 70 раз быстрее версии на чистом Питоне, с учётом того, что она использует массивы Numpy.

Одно лишь отличие от C: я использую небольшую обёртку для malloc/free, которую написал сам — cymem. Она запоминает используемые адреса памяти, и когда срабатывает сборщик мусора просто освобождает ненужную память. С тех пор, как я начал использовать эту обёртку, у меня никогда не было проблем с утечками памяти.

Промежуточный вариант писать на Cython — использовать typed memory-views, что позволяет вам работать с многомерными массивами Numpy. Однако для меня это выглядит более сложным. Обычно в своих приложениях я работаю с более простыми массивами, и предпочитаю определять свои собственные структуры данных.

Перевёл Dreadatour, текст читал %username%.