Cython将Numpy数组转为自定义结构体

技术背景

前面我们写过几篇关于Cython的文章,例如Cython计算谐振势Cython与C语言的结合Cython调用CUDA Kernel函数。Cython有着非常Pythonic的编程范式,又具备着接近于C语言的性能,因此在很多对于性能有要求的Python软件中都会使用到Cython的性能优化。Cython的基本工作流程是,先将*.pyx文件转换为*.c的C语言代码,然后再使用gcc编译成一个*.so动态链接库文件,供Python或者其他语言的代码调用。

这里我们考虑一个较为特殊的场景:将Python端常用的Numpy数组,转换为C语言的结构体,然后执行相应的任务或者计算。

Cython结构体

在Cython中可以使用ctypedef struct来定义一个可以跟C语言通用的结构体,例如,我们可以定义一个原子坐标的结构体:

python 复制代码
ctypedef struct CRD:
    double x, y, z

这是一个三维空间中的原子坐标,在Python中就相当于一个shape为(3,)的数组。如果这是C++,那我们可以通过dynamic_cast等转换方式在数组和结构体之间进行转换。而C语言和Cython则是直接通过变量类型来进行转换,以下是一个完整的Cython示例:

python 复制代码
ctypedef struct CRD:
    double x, y, z

cpdef int trans(double[:] crd):
    cdef:
        CRD* crd_s = <CRD*>&crd[0]
    print (crd_s.y)
    return 1

在这个示例中,我们输入给函数trans()的是一个shape为(3,)的数组,然后在cdef中将这个原子坐标转换为结构体的形式,最后在Cython中打印输出该原子的y坐标。我们用Cython编译这个文件之后,可以在Python中调用:

python 复制代码
In [1]: from trans import trans

In [2]: import numpy as np

In [3]: a = np.array([1,2,3], dtype=np.float64)

In [4]: trans(a)
2.0
Out[4]: 1

需要注意的是,这种转换方式并不安全,这里我们仅仅演示一下这种转换方式的使用方法。

多原子坐标结构体

其实从C语言的角度还蛮好理解的,我们通过单个结构体可以定义一个原子的空间坐标,那么用一个指针,就可以定义多个原子的空间坐标。这里我们定义多原子坐标的结构体为:

python 复制代码
ctypedef struct PATH:
    CRD crds

其实这里还是表示一个单原子的空间坐标,但是从一维指向变成了二维指向,此时我们可以通过PATH来构造一个多原子坐标的指针:

python 复制代码
ctypedef struct CRD:
    double x, y, z

ctypedef struct PATH:
    CRD crds

cpdef int trans(double[:] crd, double[:, :] path):
    cdef:
        int atoms = path.shape[0]
        int i
        CRD* crd_s = <CRD*>&crd[0]
        PATH* path_s = <PATH*>&path[0][0]
    print (crd_s.y)
    for i in range(atoms):
        print (path_s[i].crds.x, path_s[i].crds.y, path_s[i].crds.z)
    return 1

同样的,我们可以在Cython编译之后,通过Python来调用这个trans函数:

python 复制代码
from trans import trans
import numpy as np

a = np.array([1,2,3], dtype=np.float64)
b = np.arange(12).reshape((4,3)).astype(np.float64)

print (trans(a, b))

输出结果为:

bash 复制代码
2.0
0.0 1.0 2.0
3.0 4.0 5.0
6.0 7.0 8.0
9.0 10.0 11.0
1

当然,这里还是需要提醒一下,这种指针转换的方式并不安全,需要谨慎使用。

总结概要

这篇文章介绍了在Cython中定义结构体,并在Python的Numpy数组/MemoryView和自定义结构体之间进行数据转换的方法。Cython有着非常Pythonic的编程范式,又具有接近于C语言的性能,对于Python开发者而言确实是一个很棒的工具。

版权声明

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

作者ID:DechinPhy

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

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