numpy - Cython 自定义数据类型 ndarray

Question

我为我的 np.ndarrays 创建了一个数据类型:

particle_t = np.dtype([
    ('position', float, 2),
    ('momentum', float, 2),
    ('velocity', float, 2),
    ('force', float, 2),
    ('charge', int, 1),
])

根据官方的例子可以调用:

def my_func(np.ndarray[dtype, dim] particles):

但是当我尝试编译时:

def tester(np.ndarray[particle_t, ndim = 1] particles):

我收到无效类型错误。我见过的另一种用法是像 int[:] 这样的内存 View 。尝试 def tester(particle_t[:] particles): 结果:'particle_t' 不是类型标识符。

我该如何解决这个问题？

Answer 1

显然，就 Cython 而言，particle_t 不是类型，而是 Python 对象。

它类似于作为 Python 对象的 np.int32

def tester(np.ndarray[np.int32] particles):     #doesn't work!
       pass

不行，你需要使用相应的类型，即np.int32_t:

 def tester(np.ndarray[np.int32_t] particles):  #works!
      pass

但是particle_t对应的类型是什么？您需要创建一个打包结构，它将反射(reflect)您的 numpy 类型。这是一个简化版本:

#Python code:
particle_t = np.dtype([
    ('position', np.float64, 2), #It is better to specify the number of bytes exactly!
    ('charge', np.int32, 1),  #otherwise you might be surprised...
])

和对应的Cython代码:

%%cython
import numpy as np
cimport numpy as np

cdef packed struct cy_particle_t:
    np.float64_t position_x[2]
    np.int32_t   charge

def tester(np.ndarray[cy_particle_t, ndim = 1] particles):
    print(particles[0])

它不仅可以编译和加载，而且还像宣传的那样工作:

>>> t=np.zeros(2, dtype=particle_t)
>>> t[:]=42
>>> tester(t)
{'charge': 42, 'position_x': [42.0, 42.0]}