使用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倍。

总结概要

本文介绍了一下使用Cython对Python/Numpy实现的函数进行加速的一个简单案例,模型使用的是一个弹性系数全同的谐振势,然后计算总势能。从计算结果来看,使用Cython确实可以获得更接近于C语言的速度,并且编程逻辑还可以大幅的保留Python的语法。

版权声明

本文首发链接为:https://www.cnblogs.com/dechinphy/p/cython-osc.html

作者ID:DechinPhy

更多原著文章:https://www.cnblogs.com/dechinphy/

请博主喝咖啡:https://www.cnblogs.com/dechinphy/gallery/image/379634.html

参考链接

  1. https://blog.csdn.net/itcast_cn/article/details/126667915
posted @ 2024-07-24 14:52  DECHIN  阅读(40)  评论(0编辑  收藏  举报