提高Cython数组索引速度

Question

我有一个非常简单的函数需要加速。基本上，我有一个由16位数字组成的大型数组，其中有一些空洞（约10%）。我需要遍历该数组，查找连续两个0的区域，然后用前一个和下一个元素的平均值填充它们。在C语言中这只需要几毫秒，但是Python表现得要差得多。

我已经从普通的Python数组转换为了NumPy数组，然后使用Cython编译了我的代码，但是我离目标仍然相距甚远。我希望有经验的人可以看看我正在做什么并给我一些反馈。

我的普通Python代码如下：

self.rawData = numpy.fromfile(ql, numpy.uint16, 50000)
[snip]
def fixZeroes(self):
    for x in range(2,len(self.rawData)):
        if self.rawData[x] == 0 and self.rawData[x-1] == 0:
            self.rawData[x] = (self.rawData[x-2] + self.rawData[x+2]) / 2
            self.rawData[x-1] = (self.rawData[x-3] + self.rawData[x+1]) /2

我的Cython代码看起来非常相似：

import numpy as np
cimport numpy as np
DTYPE = np.uint16
ctypedef np.uint16_t DTYPE_t

@cython.boundscheck(False)
def fix_zeroes(np.ndarray[DTYPE_t, ndim=1] raw):
    assert raw.dtype == DTYPE
    cdef int len = 50000
    
    for x in range(2,len):
        if raw[x] == 0 and raw[x-1] == 0:
            raw[x] = (raw[x-2] + raw[x+2]) / 2
        raw[x-1] = (raw[x-3] + raw[x+1]) /2
    return raw

当我运行这段代码时，性能仍然比我想象的要慢得多：

开始使用Cython进行零修复

完成时间：0:00:36.983681

开始使用Python进行零修复

完成时间：0:00:41.434476

我真的认为我一定做错了什么。我看到的大多数文章都谈到了numpy和cython带来的巨大性能提升，但我只有不到10%的提升。

- Maxwell Bottiger

---

这在 C 中只需要几毫秒。如果您有 C 函数，只需直接在 Cython 中包装它，并从 Python 程序中调用即可。 - Thane Brimhall

@ThaneBrimhall 使用带有所有加速指令和类型声明的Cython应该能够非常接近C。 - Saullo G. P. Castro

@ThaneBrimhall - 说得好。我有一个C函数，但它很混乱。我需要重写这个算法的一部分，所以它被移植到Python。不过，我确实喜欢你的建议。 - Maxwell Bottiger

@SaulloCastro 是的，我知道。但是我很懒，如果我已经做完了某些东西，我通常会重复使用它。 ;) - Thane Brimhall

Answer 1

您需要声明正在用于索引raw数组的变量x:

cdef int x

你可以使用其他指令来提高性能，例如：

@cython.wraparound(False)
@cython.cdivision(True)
@cython.nonecheck(False)

- Saullo G. P. Castro

---

哇！是的，那就是问题所在。我从36.9秒变成了0.0072秒。我认为这是个胜利。 - Maxwell Bottiger

@MaxwellBottiger 很好，记住这些指令！未来你可能会面临的另一件事是将 x**2 替换为 x*x，因为前者将被转换为函数调用，类似于 pow(x, 2)。 - Saullo G. P. Castro