在 cython 中使用 numpy 掩码

Question

我正在尝试将我的一些 Python 代码转换为 Cython，当我尝试将函数定义为cdef时遇到了一些问题。

大多数问题归结为掩蔽不像在 Python 中那样工作。我想知道这是否是cdef的限制(如果我将其保留为def就可以正常工作(或者是否可以做点什么。

例如此方法

cdef func(double[:,:,:,:] arg1):
mask = arg1 > 0
...

已失败并出现编译错误：

Error compiling Cython file:
------------------------------------------------------------ ...
func (double[:,:,:,:] arg1):
mask = arg1 > 0
^
------------------------------------------------------------
cythonfile.pyx:43:20: Invalid types for '>' (double[:, :, :, :], long)

Answer 1

文档在函数定义中始终使用np.ndarray[...]，因此我会将您的函数签名更改为

cdef func(np.ndarray[np.float_t, ndim=4] arg1):

此外，您正在将float数组与long integer常量进行比较。将其更改为

mask = arg1 > 0.

将float与float进行比较.

Answer 2

double[:,:,:,:]表示法指定参数将被"解释"为类型化内存视图。这些支持许多操作，但不支持矢量化比较。

但是，将内存视图解释为函数中的 NumPy 数组非常容易：

import numpy as np
cdef func(double[:,:,:,:] arg1):
arg1arr = np.asarray(arg1)
mask = arg1arr > 0.

这甚至不需要副本，因此基本上可以"免费"在内存视图上进行np.asarray。这允许将内存视图的优点与 NumPy 数组上可能的矢量化操作相结合。

但是，对于矢量化操作，您不需要Cython，您可以在纯python函数中执行所有矢量化操作，并且仅将Cython用于"普通NumPy"函数无法实现的繁重工作。