如何用Cython对Numpy矩阵进行随机列采样以提升性能？

核心优化思路 #

原Python代码的主要性能瓶颈在于：

• 百万次循环中频繁调用Python层面的random.sample函数，带来大量函数调用开销
• 每次索引操作都要经过Python与Numpy的交互层，额外消耗资源

Cython的优化方向：

• 用C级别的随机数生成器实现无放回采样，绕过Python函数调用
• 直接操作Numpy数组的底层内存，消除Python-Numpy交互开销
• 关闭不必要的安全检查，减少循环内的额外计算

具体实现步骤 #

1. 编写Cython代码文件（sample_cols.pyx） #

import numpy as np
cimport numpy as np
from libc.stdlib cimport rand, srand
from cython cimport boundscheck, wraparound

# 关闭边界检查和负索引，最大化性能
@boundscheck(False)
@wraparound(False)
cdef void sample_columns(np.ndarray[np.float64_t, ndim=2] matrix, 
                        np.ndarray[np.intp_t, ndim=1] result_cols,
                        int n_samples):
    cdef int num_cols = matrix.shape[1]
    cdef int i, temp
    cdef np.ndarray[np.intp_t, ndim=1] indices = np.arange(num_cols, dtype=np.intp)
    
    # 用Fisher-Yates洗牌的前n步实现无放回采样，时间复杂度O(n)
    for i in range(n_samples):
        # 生成[i, num_cols-1]范围内的随机索引
        cdef int rand_idx = i + rand() % (num_cols - i)
        # 交换索引位置
        temp = indices[i]
        indices[i] = indices[rand_idx]
        indices[rand_idx] = temp
        # 保存选中的列索引
        result_cols[i] = indices[i]

def cython_sample_loop(np.ndarray[np.float64_t, ndim=2] random_matrix, 
                       int num_iterations, int n_cols_to_sample):
    cdef int iter_idx
    # 预分配采样索引数组，避免循环内重复分配内存
    cdef np.ndarray[np.intp_t, ndim=1] sampled_cols = np.empty(n_cols_to_sample, dtype=np.intp)
    
    # 初始化随机数种子（如需固定结果，替换为srand(888)这类固定值）
    srand(888)
    
    for iter_idx in range(num_iterations):
        sample_columns(random_matrix, sampled_cols, n_cols_to_sample)
        # 此处可直接使用sampled_cols获取子矩阵，示例：
        # sub_matrix = random_matrix[:, sampled_cols]
        # 若需保存结果，可提前预分配大数组存储

2. 编写编译配置文件（setup.py） #

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

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

3. 编译与调用 #

执行以下命令编译Cython代码：

python setup.py build_ext --inplace

在Python中调用优化后的函数：

import numpy as np
from sample_cols import cython_sample_loop

num_cols = 6000
random_matrix = np.random.random((20, num_cols)).astype(np.float64)
# 执行100万次采样，每次抽取5列
cython_sample_loop(random_matrix, 1_000_000, 5)

性能提升说明 #

• 用C标准库rand()替代Python的random.sample，彻底消除Python函数调用开销
• Fisher-Yates采样算法比随机选择后去重的方式更高效，尤其在采样量较小时
• 预分配数组避免循环内频繁内存申请，进一步降低开销
• 关闭边界检查和负索引，减少Cython的安全校验步骤

注意事项 #

• 若需要可复现的随机结果，将srand(888)中的数值固定即可，不要用动态时间种子
• 确保输入的Numpy矩阵类型与Cython代码中声明的np.float64_t一致，否则会触发类型转换开销
• 如果矩阵是其他数据类型（如float32），需修改Cython代码中的类型声明为对应类型

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