Cython 101

注：本文同步也发表在我的独立博客中。

从网上照抄的一些 LBM 算例基本都是在规整的矩形区域内的计算，用简单的 Numpy 非常适合，但当需要计算一些稍微更复杂的计算域时，纯 Numpy 的计算逻辑就有点复杂了，这时最简单的办法还是循环计算节点，这时候就不得不请出 Cython 了。

对我这种级别的人来说，Cython 基本上可以理解成用 Python 的语法来写 C，与 Python 的主要区别在于可以定义变量的类型，以及需要编译成 Python 库才能运行。这篇文章简单记录一下使用 Cython 的一些入门知识。

使用 Cython，先需要建立一个 `pyx` 格式的源文件，以计算宏观密度的函数为例，建立 `rho.pyx` 文件，并写入如下的函数：

def calc_rho_notype(f):
    rho = np.empty_like(f[:,:,0])
    for i in xrange(f.shape[0]):
        for j in xrange(f.shape[1]):
            for k in xrange(f.shape[2]):
                rho[i,j] += f[i,j,k]
    return rho

由于使用到了 Numpy，需要在程序前面加上：

import numpy as np
cimport numpy as np

注意这里除了 `import` 外，还需要再 `cimport` 一次才行。

编译这个源代码有多种办法，这里讲最基础的一种，在同目录下建立 `setup.py` 文件，文件内容为：

from distutils.core import setup
from distutils.extension import Extension
from Cython.Distutils import build_ext
import numpy

ext_module = Extension(
    "rho",
    ["rho.pyx"],
    include_dirs=[numpy.get_include()],
)

setup(
    name = 'Cython Test',
    cmdclass = {'build_ext': build_ext},
    ext_modules = [ext_module],
)

注意这里跟 Numpy 的地方有两处，一处是开头的导入，另一处是告诉编译器需要包括相关的 h 文件。

有了 setup.py 之后，在 cmd 中执行：`python setup.py build_ext --inplace`，编译不出错的话就会生成 `rho.pyd` 文件，这个就可以在主函数中导入了。

但是，这样直接使用 Python 代码生成的效率仍然很低，因为里面的变量仍是 Python 对象，作为对比，可以再写一个 Cython 函数：

def calc_rho_type(np.ndarray[double,ndim=3] f):
    cdef np.ndarray[double,ndim=2] rho = np.empty_like(f[:,:,0])
    cdef int i,j,k
    for i in xrange(f.shape[0]):
        for j in xrange(f.shape[1]):
            for k in xrange(f.shape[2]):
                rho[i,j] += f[i,j,k]
    return rho

对比的结果在后文。

还有一个重要的措施可以提高运行速率，即让它并行化！Cython 的并行库用的是 Openmp，使用也非常简单，Cython 函数如下：

def calc_rho_openmp(np.ndarray[double,ndim=3] f):
    cdef np.ndarray[double,ndim=2] rho = np.empty_like(f[:,:,0])
    cdef int i,j,k
    with nogil:
        for i in prange(f.shape[0]):
            for j in xrange(f.shape[1]):
                for k in xrange(f.shape[2]):
                    rho[i,j] += f[i,j,k]
    return rho

注意其中为了并行需要释放 GIL，并在外层循环中把 `xrange` 写成了 `prange`。这样标记表明这个循环的结果与执行顺序无关，可以并行，为了使用它，需要在开关添加：

from cython.parallel import prange

并且 setup.py 也要稍作修改，在 `ext_module` 中增加两个参数：

extra_compile_args=['-fopenmp'],
extra_link_args=['-fopenmp']

其它流程都是一样的。

下面用 iPython 自带的 timeit 工具看看各种方法的效率：

In [1]: import numpy as np

In [2]: import numexpr as ne

In [3]: import rho

In [4]: f = np.random.random([1000,1000,9])

In [5]: %timeit np.sum(f,axis=2)
10 loops, best of 3: 31.3 ms per loop

In [6]: %timeit ne.evaluate("sum(f,axis=2)")
10 loops, best of 3: 122 ms per loop

In [7]: %timeit rho.calc_rho_notype(f)
1 loops, best of 3: 6.52 s per loop

In [8]: %timeit rho.calc_rho_type(f)
10 loops, best of 3: 70.3 ms per loop

In [9]: %timeit rho.calc_rho_openmp(f)
10 loops, best of 3: 35 ms per loop

这里为了比较，还用了之前提到过的 Numexpr 包。可以看到，在不声明类型的情况下，这种嵌套循环的程序效率还是非常低，而添加类型就让执行速度提升了接近100倍，使用 Openmp 之后，CPU 迅速升到 90% 以上，但带来的速度提升是50%（4核心）。值得注意的是，在这个简单的例子中， Numexpr 包以及 Cython 并行的函数两者的计算速度都输给了 Numpy 原生的数组求和函数，一方面这说明 Numpy 原生函数确实是经过了高度优化，平时还是应当尽可能使用它们，但另一方面，由于这是逻辑非常简单的计算，而在更复杂的计算中，例如在将 f 求和之前，先作 `f**1.5` 的操作，结果就不同了：

In [5]: %timeit np.sum(f**1.5,axis=2)
1 loops, best of 3: 891 ms per loop

In [6]: %timeit ne.evaluate("sum(f,axis=2)")
10 loops, best of 3: 117 ms per loop

In [7]: %timeit rho.calc_rho_notype(f)
1 loops, best of 3: 8.19 s per loop

In [8]: %timeit rho.calc_rho_type(f)
1 loops, best of 3: 1.45 s per loop

In [9]: %timeit rho.calc_rho_openmp(f)
1 loops, best of 3: 382 ms per loop

可以看出，这时候，Numexpr 以及 并行 Cython 都开始出现了明显的优势。考虑到 Numexpr 极其简单的使用方法，还是推荐在绝大多数逻辑不复杂的时候，尽量使用它。

最后值得一提的是，在我稍复杂的 LBM 程序中，核心计算函数都采用并行 Cython 比 Numexpr 快 30% 左右。