从网上照抄的一些 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% 左右。