使用Cython加速Numpy对象数组按行求和的技术求助

嘿，我来帮你搞定这个numpy对象数组按行求和的效率难题！确实，对象数组的原生求和因为Python层的额外开销慢得离谱，咱们得从内存布局和Cython的正确用法入手解决。

先理清楚为什么原生方法差距这么大：

• 普通numpy数组a是连续的内存块，sum(1)直接在C层遍历连续数据，效率极高；
• 而对象数组a_obj本质是存了一堆指向numpy子数组的指针，每个子数组的内存不连续，原生sum()要逐个调用子数组的求和方法，每次都要经过Python对象的操作，开销直接拉满，慢250倍完全合理。

你之前的两个Cython尝试踩了典型的坑：

1. 直接遍历a[i][j]：本质还是通过Python的属性访问和__getitem__，根本没绕过Python层，所以速度和原生遍历差不多，完全没提升；
2. 直接转double**：对象数组的data里存的是指向numpy数组对象的指针，不是直接指向double数值的内存！所以你读出来的是对象的内存地址，自然是像5e-324这种无效值。

核心思路：拿到每个子numpy数组的底层连续数据指针，完全在C层遍历求和，彻底避开Python对象的操作开销。

1. 编写Cython代码 #

把下面的代码保存为row_sum_obj.pyx：

import numpy as np
cimport numpy as np
from cpython cimport PyObject

def row_sum_obj(np.ndarray[PyObject, ndim=1] a_obj):
    cdef:
        int n = a_obj.shape[0]
        # 假设所有子数组的长度一致（和你的测试数据匹配）
        int row_count = (<np.ndarray> a_obj[0]).shape[0]
        # 初始化结果数组，存每行的和
        np.ndarray[np.double_t, ndim=1] result = np.zeros(row_count, dtype=np.double)
        np.ndarray sub_arr
        double *sub_data_ptr
        int i, j
    
    for i in range(n):
        # 取出当前子数组（转成numpy数组对象）
        sub_arr = <np.ndarray> a_obj[i]
        # 获取子数组的底层连续double数据指针
        sub_data_ptr = <double*> sub_arr.data
        # C层循环累加求和
        for j in range(row_count):
            result[j] += sub_data_ptr[j]
    
    return result

2. 编译Cython代码 #

创建setup.py文件：

from setuptools import setup
from Cython.Build import cythonize
import numpy as np

setup(
    ext_modules=cythonize("row_sum_obj.pyx"),
    include_dirs=[np.get_include()]
)

然后执行编译命令：

python setup.py build_ext --inplace

3. 测试效果 #

编译完成后，用下面的代码验证正确性和速度：

import numpy as np
from row_sum_obj import row_sum_obj

# 构造测试数据
n = 10**5
a = np.random.rand(10, n)
a_obj = np.empty(n, dtype='O')
for i in range(n):
    a_obj[i] = a[:, i]

# 先验证结果和原生方法一致
assert np.allclose(row_sum_obj(a_obj), a.sum(1))

# 测试速度
%timeit row_sum_obj(a_obj)

这个实现的速度应该和a.sum(1)非常接近，因为全程在C层操作连续内存，没有Python循环的额外开销。

如果你的子数组可能不是连续内存（比如是切片得到的非连续数组），可以在Cython里先转成连续数组，避免访问非连续内存的性能损耗：

sub_arr = np.ascontiguousarray(<np.ndarray> a_obj[i])

不过如果你的子数组都是像测试数据里那样连续的（a[:,i]是连续的），这一步可以省略。

内容的提问来源于stack exchange，提问作者Ted Petrou