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

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

技术背景

前面我们写过几篇关于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语言通用的结构体，例如，我们可以定义一个原子坐标的结构体：

ctypedef struct CRD:
    double x, y, z
  • 1.
  • 2.

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

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
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.

在这个示例中，我们输入给函数trans()的是一个shape为(3,)的数组，然后在cdef中将这个原子坐标转换为结构体的形式，最后在Cython中打印输出该原子的y坐标。我们用Cython编译这个文件之后，可以在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
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.
  • 9.

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

多原子坐标结构体

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

ctypedef struct PATH:
    CRD crds
  • 1.
  • 2.

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

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
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.
  • 9.
  • 10.
  • 11.
  • 12.
  • 13.
  • 14.
  • 15.
  • 16.

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

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))
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.

输出结果为：

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
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.

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

总结概要

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


作者ID：DechinPhy