Cython初体验

python

Cython初体验

写一些仿真的程序，用Python很快速地就可以把项目的框架搭建起来。我自己实践的基本流程是建分支写对应功能，用Jupyter调试新写的.py模块，调试好了直接把Jupyter中的代码整理到tests目录下作为单元测试，在每一次commit前都保证单元测试是通过的，这样保证每一个节点的程序拿出来都是可以运行的。但是因为Python是动态语言，因此在计算密集的地方比Fortran/C/C++慢上100-300倍是很正常的。但是

Programmers waste enormous amounts of time thinking about, or worrying about, the speed of noncritical parts of their programs, and these attempts at efficiency actually have a strong negative impact when debugging and maintenance are considered. We should forget about small efficiencies, say about 97% of the time: premature optimization is the root of all evil. Yet we should not pass up our opportunities in that critical 3%.

整个程序真正耗时的核心是3%左右的部分，程序的大部分模块我们是不需要花费精力考虑优化的。所以一个兼顾开发效率和计算效率的工作方式就是，先用Python开发出大部分功能和框架，耗时的部分用C/C++重写再嵌入到Python的框架中。实现这个流程的关键就是程序的各个模块之间是弱耦合的。软件设计中的一个基本原则就是单一职责原则（Single Responsibility Principle, SRP），它指出

一个模块应该只有一个单一的功能。

在邹欣老师的《构建之法》中提到过《敏捷软件开发：原则、模式与实践》中举过的一个例子：

一个处理正方形的模块有两个功能：1. 计算面积，2. 画出这个正方形。

这个设计让一个模块负责两个不同的职责：进行几何计算（与显示图形无关）和在图形界面绘出正方形。如果是一个集合计算的程序需要使用这个模块，那么它要同时包括图形显示的部分（因为是在同一个模块中），这是一种浪费，也引入了不必要的依赖（因为图形显示和具体的图形底层的实现相关），妨碍了可移植性。另外，几何计算需求的改变和图形显示需求的改变都会导致这个模块发生变化，增加错误发生的风险。实际上科学计算也是差不多的，假如我们所有的仿真流程都塞到了一个或几个模块里，每个函数都需要调用多个对象，实现了多个功能，因为彼此之间相互依赖，这给用C/C++对某些核心部分重写带来了很大的困难。虽然不同设计的程序都可以实现对某个对象的仿真，但是不同的设计会让整个系统看起来有很大差别。就像庖丁解牛一样，虽然牛都是一样的，但是不同的刀法会让切出来的牛看起来很不一样.. 发生的事情都是一样的，关键是用哪一套图像解构再重构发生的事情，来帮助我们理顺事情发生的逻辑并做出一些预测。理解不是记忆，而是解构和建构，所以一些时候学习了很多东西但是始终没有什么想法，主要是缺乏了在实践中用自己的知识重构自己所面对的现象的过程，离开实践的认识是不可能的。

一般情况下，当我们执行Python代码时，首先Python解释器会进行词法分析，从源代码生成字节码（也就是__pycache_目录中.pyc的文件）。然后，Python虚拟机（解释器的一部分）读取并执行这些字节码。因为Python中一切皆对象，连简单的int类型实际上也是PyIntObject对象，这会导致代码运行的过程包含了远超过它看起来做的事情的环节，比如创建和销毁对象、类型检查等等.. 而对于静态语言来说，当我们把源代码编译成机器码之后，可以被操作系统和硬件直接执行，而不需要通过虚拟机的加入，避免了运行时解释的开销。所以提升Python程序性能的基本想法，就是把部分耗时的代码先提前进行编译，不再让他们和解释器打交道。目前有三个典型的方案可以部分实现这个需求：Numba、Pybind11和Cython。

Numba最简单，直接在函数上加一个装饰器numba.jit(python)就可以了，Numba会将Python的字节码转换成LLVM（Low Level Virtual Machine）中间表示，然后再将其编译成机器码。一旦函数被编译了，后续的函数调用将直接运行机器代码，绕过Python虚拟机。但是Numba有些问题，虽然说是可以直接在函数上加装饰器而无需改变原有代码，但封装的程度太高了，基本没法定制，而且并不是所有的Python代码都可以被Numba优化，有时候需求复杂了不得不写出些和源程序差异很大的代码。另外，由于JIT是即时编译，当装饰的函数多了以后，每次import一遍库都需要从头到尾编译一遍，这也引入了些不便。最后整个项目即不像Python也不像C。Numba比较适合简单的原型开发时做一些调试。

Pybind11主要是用来将C++代码封装为Python模块，虽然据说效果非常好，但是似乎一个更合适的场景是已经开发了一个完整的C++模块，然后想去提供一个Python的接口，在Python中调用这个C++模块。Cython介乎两者之间，可以将Python代码编译成C的代码，从而获得性能提升，也支持C/C++库的调用。Cython前两天发布了3.0版本，官方文档上这样介绍，

Cython 项目通过源代码编译器将 Python 代码转换为等效的 C 代码来解决这个问题。此代码在 CPython 运行时环境中执行，但是却以编译后的 C程序那般速度执行，并且能够直接调用 C语言库。同时，它保留了 Python 源代码的原始接口，这使得它可以直接使用Python语言代码。这些双重特性使 Cython 的这两个主要使用场景成为可能：使用快速二进制模块来扩展 CPython 解释器，以及将 Python 代码与外部 C 库连接。

与此同时 Cython 可以编译（大多数）常规 Python 代码，而且生成的 C 代码通常可以从 Python 和 C 类型的任意静态类型声明中获得主要（并且有时很惊人）的速度上的提升。这些允许 Cython 将 C 语义分配给代码的一部分，并将它们转换为非常高效率的 C 代码。因此，类型声明可用于两个目的：将代码段从动态 Python 语义转换为静态和快速 C 语义，还用于直接操作外部库中定义的类型。因此，Cython 将这两个世界合并为一种非常广泛适用的编程语言。

上周在自己的蒙特卡洛程序上试用了一下，感觉效果还挺好的，这个Blog是记录一下在自己的程序中使用Cython做优化的经历。

性能分析

进行性能分析是代码优化前的必要步骤，首先需要定位程序的热点，可能整个程序运行100s，其中某一个函数占了90s，剩下的所有其他内容占据了10s，这时候大部分情况下我们只优化这个占据了90s的函数就足够了。Python官方提供的性能分析组合是cProfile模块+pstats模块，其中cProfile负责统计函数或脚本执行中各部分消耗的时间和调用次数并生成统计信息文件，pstats负责处理分析得到的统计信息。cProfile的基本使用非常简单，只要第一个参数放上去要运行统计的函数，第二个参数指定保存的文件名就可以了。

1
2
3
4
5
import cProfile
import pstats
import some_Monte_Carlo_module
......
cProfile.run('MC.run()', 'simulation.prof')

之后，我们用pstats模块分析这个.prof结果，其中strip_dirs()是指在结果中只保留模块名，忽略掉模块的绝对路径；sort_stats()指定按照哪一个指标来排序，print_stats(10)指定仅保留前10个结果。

1
pstats.Stats('simulation.prof').strip_dirs().sort_stats('cumtime').print_stats(10)

得到的结果是这样的

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
Mon Sep 11 13:51:17 2023    origin.prof

         12463639 function calls (11974174 primitive calls) in 13.530 seconds

   Ordered by: cumulative time
   List reduced from 62 to 10 due to restriction <10>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000   13.530   13.530 {built-in method builtins.exec}
        1    0.000    0.000   13.530   13.530 <string>:1(<module>)
        1    0.063    0.063   13.530   13.530 simulation.py:131(run)
    10000    0.872    0.000   13.125    0.001 simulation.py:70(track)
   163155    0.175    0.000    9.094    0.000 surface.py:216(calculate_intersection)
   163155    0.997    0.000    8.919    0.000 surface.py:80(calculate_intersection_between_line_segment_and_surface)
978930/489465    1.100    0.000    7.591    0.000 {built-in method numpy.core._multiarray_umath.implement_array_function}
   163155    0.090    0.000    7.284    0.000 <__array_function__ internals>:177(isclose)
   163155    0.588    0.000    6.962    0.000 numeric.py:2278(isclose)
   163155    0.969    0.000    3.294    0.000 numeric.py:2359(within_tol)

下面这几列分别是，

  • ncalls，调函数调用次数

  • tottime，函数内部消耗时间，不包括这个函数调用其他函数的时间

  • percall，= tottime / ncalls

  • cumtime，这个函数消耗的总时间

  • percall, = cumtime / ncalls

  • filename:lineno(function)，函数所在的文件名:所在行数(函数名)

从上面的结果可以看到，这次仿真一共消耗了13.530s，在整个run函数中，最耗时的部分就是调用track函数跟踪粒子的过程，track函数消耗了13.125s。在track函数中，最耗时的函数就是calculate_intersection函数，它主要负责调用calculate_intersection_between_line_segment_and_surface函数来计算线段和某个面的交点，这个函数占据了8.919s，计算交点主要用来处理边界和粒子的散射过程。所以要用Cython提高仿真的效率，最关键的就是先把这个函数给优化了。这个函数的逻辑非常简单，检查给定的线段是否与这个面相交，如果相交，则返回线和面的交点。函数输入线段的两个端点、平面上的一个点、以及平面的法向向量。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
def calculate_intersection_between_line_segment_and_surface(line_point_A: NDArray, line_point_B: NDArray, surface_point: NDArray, surface_normal: NDArray) -> Optional[NDArray]:
    """检查给定的线段是否与此面相交，并返回交点（如果有）。

    参数:
        line_point_A (NDArray): 线段的起点。
        line_point_B (NDArray): 线段的终点。
        surface_point: (NDArray): 面上的一点。
        surface_normal (NDArray): 面的法线向量

    返回:
        Optional[NDArray]: 如果相交，返回交点；否则返回None。
    """

    # 计算线段方向
    line_direction = line_point_B - line_point_A

    # 计算线段和面的交点
    denominator = np.dot(surface_normal, line_direction)

    # 检查线段和面是否平行
    if np.isclose(denominator, 0):
        return None

    t = np.dot(surface_normal, (surface_point - line_point_A)) / denominator
  
    # 检查交点是否在给定的线段上
    if 0 <= t <= 1:
        intersection_point = line_point_A + t * line_direction
        return intersection_point

    return None

构建Cython代码

https://www.kancloud.cn/apachecn/cython-doc-zh/1944771

Cython代码写在.pyx文件中，与Python代码不同，Cython代码必须编译。这包含两个阶段

  • .pyx文件由Cython编译为.c文件，它含有Python扩展文件的代码。
  • .c文件由C编译器编译成Python的动态链接库（Windows平台生成.pyd文件，Linux平台生成.so文件），可以直接被import进入到一个Python程序中。Distuils或setuptools负责这一部分。

假如我们我们之前的Python代码放到了Sim目录下，我们现在在Sim目录下新建了一个Cython/surface.pyx文件，里面放上了用Cython重写过的函数，现在想把它打包成一个库导入供其他模块调用，可以在项目总目录下新建一个setup.py文件。setup.py是Python中用于配置和构建包的脚本。它用来指定包的元数据（如名字、版本、作者等）、依赖关系、构建指令和其他设置。对于一个包，我们在终端中执行python setup.py install就可以把这个包装到环境中了。在setup.py中，setup中的ext_modules参数用于指定要构建和安装的C扩展模块，这些扩展模块通常是用C或C++编写的，并可以从Python代码中调用。在使用Cython时，我们可以用cythonize函数来创建也给Extension对象列表，这样Cython就可以处理.pyx文件的编译，把它们转换成C代码。之后， setup.py 脚本调用 Python 的 setuptools 或 distutils 模块来继续完成转换后的C代码的编译和链接，最后变成ext_modules成为可供Python调用的动态链接库。比如对于下面的setup.py，我们指定了一个C模块Sim.Cython.surface，这个模块的源文件是Sim/Cython/surface.pyx，需要numpy的头文件numpy.get_include()。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
from setuptools import setup, Extension, find_packages

from Cython.Build import cythonize
import numpy

cython_modules = [
    Extension(
        'Sim.Cython.surface',
        sources=['Sim/Cython/surface.pyx'],
        language='c',
        include_dirs=[numpy.get_include()],
    )
]

setup(
    name='Sim
    version='1.0',
    description='xxxx simulation',
    packages=find_packages(),
    install_requires=['numpy', 'nptyping', 'Cython'],
    extras_require={
        "test": ["pytest"],
        "plotting": ['plotly', 'kaleido'],
    },
    ext_modules=cythonize(cython_modules),
)

我们执行python setup.py build_ext --inplace命令，就可以完成这个Cython模块的构建，整个过程可以分为

  1. Cython将.pyx文件转换成.c文件。
  2. C编译器将.c文件编译成对象文件。
  3. 链接器将对象文件链接成一个动态链接库。

其中--inplace参数指定把生成的.pyd文件拷贝到与.pyx源文件相同的目录中，使得我们可以直接从源代码目录中导入这个新构建的模块，而不需要手动移动文件或者修改Python路径。

对于我们上面的Cython模块，我们运行setup.py时，它会创建一个build文件夹📂来存放构建过程中生成的各种文件和最终的输出文件。在我的电脑上，生成了lib.win-amd64-cpython-38和temp.win-amd64-cpython-38两个目录，其中lib目录下储存了构建过程的最终输出，即Cython的扩展编译模块，也就是.pyd文件（在Windows上）或.so文件（在Unix-like系统上）。temp目录包含了构建过程中生成的临时文件，比如.obj对象文件，库文件.lib文件，包含符号信息的.exp文件等等，他们没什么用。构建完成了之后，我们就可以在项目的Python模块中导入这个模块了from .Cython.surface import calculate_intersection_between_line_segment_and_surface。

Cython代码

Cython编译的代码，有些是纯C的，只能在Cython中调用，而有些是可以直接在Python中调用的，另一方面，可以要大致理解Cython优化代码的图像，先要简单理解一下CPython。CPython就是Python语言的最主流实现，也就是用C语言写的Python。在CPython里，每一个Python对象都是一个C的PyObject结构体，其中ob_refcnt是引用计数，Python的垃圾回收机制靠它来实现，当一个对象的引用计数为0时，这个对象就被回收。ob_type是类型指针，指向对象的类型对象。

1
2
3
4
5
typedef struct _object {
    _PyObject_HEAD_EXTRA
    Py_ssize_t ob_refcnt;
    struct _typeobject *ob_type;
} PyObject;

所以在CPython解释器中本质上所有的Python对象实际上都是C语言层面的结构体。这意味着Python的数据类型和对象背后有对应的C数据结构来实现和存储数据。在Cython中，纯Python的代码也可以编译，就是把它们还原成复杂的结构体嘛。同时，我们也可以通过一些类型定义，跳过这些结构体，直接访问和操作C级别的底层数据，以实现更高效的代码，自然Cython也允许混合Python和C代码，但是显然我们应该尽可能地跳过这些结构体，直接在C层面运行代码，以避免Python层面的开销来提高代码的运行效率。在Cython中，我们使用cdef来定义C层面的变量或函数，他们只能在Cython模块内部访问，而无法被外部模块访问。我们可以通过一个小例子来查看这种类型的巨大差别，定义一个cython_test.pyx模块，

1
2
3
4
5
6
7
8
9
# cython_test.pyx
from sys import getsizeof

def check_sizes():
    cdef int c_int = 1
    python_int = 1

    print("Size of Python int object: ", getsizeof(python_int), "bytes")
    print("Size of C int object: ", sizeof(c_int), "bytes")

当我们用Cython编译并调用这个函数时，得到的结果将是

1
2
Size of Python int object:  28 bytes
Size of C int object:  4 bytes

一个Python的int对象是28字节，而一个C的int类型只有4个字节，这会带来巨大的性能差异，所以尽可能要在C的层面完成Cython的模块，只在输出输出层面提供一个Python类型的封装。在Cython中另一个常用的定义是cpdef，cpdef的想法就是能转换成C底层变量的，Cython就会转换，以提高性能。如果不能的，就回退到使用Python的C-API来处理这些对象。所以当我们采用cpdef重新定义一个函数时，Cython会尽可能地把内部流程直接通过底层的C类型实现，同时也会在输入输出上自动实现一些Python对象和C类型之间的转换。

下面一个采用Cython进行初步优化的calculate_intersection_between_line_segment_and_surface函数，其中：

  • # cython: language_level=3：指定Python语言的版本为3

  • cimport用于在Cython中导入C语言的库或者其他Cython模块，cimport numpy就是导入了numpy给Cython预留的C级别接口

  • @cython.boundscheck(False)禁用了数组边界检查。在Python中，当尝试访问一个数组或列表的一个超出其边界的索引时，Python会引发一个IndexError异常。但是这种边界检查在运行时会增加额外的开销。

  • @cython.wraparound(False)禁用了Python的负索引包装。在Python中可以使用负索引来从数组的末尾开始计数，例如arr[-1]将返回数组的最后一个元素，但这也会带来额外的开销。

  • cnp.ndarray[cnp.float64_t, ndim=1]指定了参数的类型，Cython可以根据提供的类型信息进行更多的优化，可以在C级别直接访问numpy数组的内存，而不需要通过Python对象层。

  • cdef在C层面声明了变量类型，允许Cython生成直接的C代码来处理这些变量。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
# cython: language_level=3
import numpy as np
cimport numpy as cnp
cimport cython

from libc.math cimport fabs, NAN

@cython.boundscheck(False)
@cython.wraparound(False)
cpdef cnp.ndarray[cnp.float64_t, ndim=1] calculate_intersection_between_line_segment_and_surface(
        cnp.ndarray[cnp.float64_t, ndim=1] line_point_A, 
        cnp.ndarray[cnp.float64_t, ndim=1] line_point_B, 
        cnp.ndarray[cnp.float64_t, ndim=1] surface_point, 
        cnp.ndarray[cnp.float64_t, ndim=1] surface_normal):
    """检查给定的线段是否与此面相交，并返回交点（如果有）。
    
    参数:
        line_point_A (cnp.ndarray[cnp.float64_t, ndim=1]): 线段的起点。
        line_point_B (cnp.ndarray[cnp.float64_t, ndim=1]): 线段的终点。
        surface_point: (cnp.ndarray[cnp.float64_t, ndim=1]): 面上的一点。
        surface_normal (cnp.ndarray[cnp.float64_t, ndim=1]): 面的法线向量。

    返回:
        cnp.ndarray[cnp.float64_t, ndim=1]: 如果相交，返回交点；否则返回None。
    """
    cdef:
        cnp.ndarray[cnp.float64_t, ndim=1] line_direction = np.zeros(3, dtype=np.float64)
        cnp.float64_t denominator, t
        cnp.ndarray[cnp.float64_t, ndim=1] intersection_point = np.zeros(3, dtype=np.float64)

    # 计算线段方向
    line_direction[0] = line_point_B[0] - line_point_A[0]
    line_direction[1] = line_point_B[1] - line_point_A[1]
    line_direction[2] = line_point_B[2] - line_point_A[2]

    # 计算线段和面的交点
    denominator = surface_normal[0]*line_direction[0] + surface_normal[1]*line_direction[1] + surface_normal[2]*line_direction[2]

    # 检查线段和面是否平行
    if fabs(denominator) <= 1e-12:
        intersection_point[0] = NAN
        intersection_point[1] = NAN
        intersection_point[2] = NAN
        return intersection_point

    t = (surface_normal[0]*(surface_point[0] - line_point_A[0]) + 
         surface_normal[1]*(surface_point[1] - line_point_A[1]) + 
         surface_normal[2]*(surface_point[2] - line_point_A[2])) / denominator

    # 检查交点是否在给定的线段上
    if 0.0 <= t <= 1.0:
        intersection_point[0] = line_point_A[0] + t * line_direction[0]
        intersection_point[1] = line_point_A[1] + t * line_direction[1]
        intersection_point[2] = line_point_A[2] + t * line_direction[2]
        return intersection_point

    intersection_point[0] = NAN
    intersection_point[1] = NAN
    intersection_point[2] = NAN
    return intersection_point

我们可以比较一下优化前后的运行效率，这里把Numba装饰的版本也一并放上比较。

  • Python实现：42.2 µs ± 951 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
  • Numba装饰：3.66 µs ± 49.7 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
  • Cython编译：3.55 µs ± 23.7 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

有超过10倍的性能提升，重新运行我们的MC.run()进行模拟，这回可以发现，仅仅优化了一个函数，仿真程序的计算效率提高了4倍，耗时的主要部分也变成了抽样声子属性的函数，可以接下来进行进一步的优化。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
Mon Sep 11 20:10:07 2023    new2.prof

         1983676 function calls in 2.946 seconds

   Ordered by: cumulative time
   List reduced from 26 to 10 due to restriction <10>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000    2.946    2.946 {built-in method builtins.exec}
        1    0.000    0.000    2.946    2.946 <string>:1(<module>)
        1    0.049    0.049    2.946    2.946 simulation.py:131(run)
    10000    0.752    0.000    2.705    0.000 simulation.py:70(track)
   121608    0.051    0.000    0.663    0.000 material.py:37(sample_phonon_direction)
   121608    0.430    0.000    0.611    0.000 material.py:10(sample_direction_in_sphere)
   161840    0.073    0.000    0.345    0.000 surface.py:216(calculate_intersection)
   323816    0.343    0.000    0.343    0.000 geometry.py:80(is_inside)
    30368    0.017    0.000    0.300    0.000 surface.py:264(handle_phonon)
    20368    0.099    0.000    0.281    0.000 surface.py:160(adiabatic_strategy)