使用Cython加速谐振势的计算

技术背景

Python的计算性能在很多计算密集型任务中都是被诟病的对象，所以Python开发者经常要寻求一些加速的手段。这里暂且不讨论GPU加速和多进程的模式，因为这些方法更依赖于硬件加速，我们只考虑软件加速。之前的一篇博客介绍过numba的计算加速，其实我个人是比较喜欢这种可以最大化保留Python代码的加速方式的，但是numba在很多测试场景中的加速效果是比较有限的（有场景依赖）。还有一种方法是使用C/C++语言进行编程，然后编译成动态链接库给Python去调用。这种形式无疑能够给性能带来较大的提升，但基本放弃了Python编程语言的便捷性。除了numba和动态链接的方案，我们还有pypy（基于CPython）和Cython的方案，本文介绍的是Cython的加速方案。

Python谐振势计算

我们考虑一个最简单的计算场景：谐振势。谐振势的计算比较简单：

\[P=\sum_{i,j}k_{i,j}\textbf{r}_{i,j}^2,i<j \]

如果用numpy来实现，我们甚至可以一行代码搞定：

P = lambda r:np.square(np.triu(np.linalg.norm(r[:, None]-r[None], axis=-1))).sum()

简单总结一个思路就是，利用numpy.ndarray的自动广播机制计算两两间距，但是这里因为每一个距离都被计算了两次，因此我们取一个上三角再做求和，得到的才是真正的谐振势的结果。当然，这种计算写起来容易，但是性能上永远受限于\(O(N^2)\)的复杂度。因此，还可以考虑用循环来替代广播机制，以缩减内存空间的占用：

def loop_osc(r):
    pot = 0.0
    for i in range(r.shape[0]-1):
        pot += np.sum(np.linalg.norm(r[i][None]-r[i+1:], axis=-1) ** 2)
    return pot

Cython谐振势计算

Cython的语法跟Python是非常类似的，文件以*.pyx的形式存储，在pyx文件中可以直接引用和使用python的普通函数，也可以通过cimport来引用cython版本的函数。Cython的函数跟Python函数的写法上区别就是函数定义def变成了cpdef，变量定义使用的是cdef进行声明。如果不使用cdef声明，表示一个普通的Python变量（在Cython函数中可以调用/使用普通Python的函数和变量），会受到GIL的约束，性能只比普通的Python好一点点。我们参照上面这个循环的案例写一个Cython的版本：

# cyosc.pyx
cimport cython

@cython.boundscheck(False)
@cython.wraparound(False)
cpdef double cy_osc(double[:, :] r):
    cdef:
        int atoms = r.shape[0]
        double pot = 0
        int i, j
    for i in range(atoms-1):
        for j in range(i+1, atoms):
            pot += (r[i][0]-r[j][0])**2
            pot += (r[i][1]-r[j][1])**2
            pot += (r[i][2]-r[j][2])**2
    return pot

这里两个装饰器仅仅是取消了边界检查，以免导致索引溢出的问题。虽然确实优化了一点点性能，但是相比于总体时间来说还是微乎其微的。保存成文件cyosc.pyx之后，需要使用如下指令编译一次：

$ cythonize -i cyosc.pyx

编译正确执行后，会在当前路径下生成一个c文件和一个so动态链接库文件：

-rw-r--r-- 1 root root  994546 Jul 24 11:09 cyosc.c
-rwxr-xr-x 1 root root  911296 Jul 24 11:09 cyosc.cpython-37m-x86_64-linux-gnu.so*
-rw-r--r-- 1 root root     429 Jul 24 11:09 cyosc.pyx

完整测试案例

把上面这几个函数组合起来进行测试，我们生成一个随机的初始坐标，然后分别用三个函数进行计算，并统计耗时：

import numpy as np

def np_osc(r):
    return np.square(np.triu(np.linalg.norm(r[:, None]-r[None], axis=-1))).sum()

def loop_osc(r):
    pot = 0.0
    for i in range(r.shape[0]-1):
        pot += np.sum(np.linalg.norm(r[i][None]-r[i+1:], axis=-1) ** 2)
    return pot

if __name__ == '__main__':
    import time
    from cyosc import cy_osc
    np.random.seed(1)
    atoms=20000
    r=np.random.random((atoms, 3))
    time0 = time.time()
    potential1 = np_osc(r)
    time1 = time.time()
    potential2 = loop_osc(r)
    time2 = time.time()
    potential3 = cy_osc(r)
    time3 = time.time()
    print ('The potential by numpy is: {} in time {} sec.'.format(potential1,
                                                                  time1-time0))
    print ('The potential by loop is: {} in time {} sec.'.format(potential2,
                                                                  time2-time1))
    print ('The potential by cython is: {} in time {} sec.'.format(potential3,
                                                                   time3-time2))

得到的结果为：

The potential by numpy is: 100149362.95950225 in time 13.057415246963501 sec.
The potential by loop is: 100149362.95950234 in time 4.418238162994385 sec.
The potential by cython is: 100149362.95931925 in time 0.24606728553771973 sec.

这说明，Cython在这个谐振势计算的案例中加速了大约60倍（相比于普通numpy的广播计算），即使是相比于同样使用了循环的numpy计算，也优化了接近20倍。