Cython:双精度或数组作为输入

Question

让我们考虑下面这个简单的例子:

def add(double a, double b):
cdef double c
c = a + b
return c

我如何定义它，使输入a和b(以及输出c)可能是double或np.array?

非常感谢!

Answer 1

对于1D NumPy数组:

def add(np.ndarray[np.float64_t, ndim=1] a, np.ndarray[np.float64_t, ndim=1] b):
cdef np.ndarray [np.float64_t, ndim=1] c
c = a + b
return c

对于2D NumPy数组:

def add(np.ndarray[np.float64_t, ndim=2] a, np.ndarray[np.float64_t, ndim=2] b):
cdef np.ndarray [np.float64_t, ndim=2] c
c = a + b
return c

可以看到ndim是维数。np.float64_t是数组元素np.float64的数据类型。np.ndarray将数据类型定义为NumPy数组。

对于double的数组，您可以查看memoryviews，但这可能不是您想要的。它用于快速索引，而不是算术运算。