如何在Cython中向二维数组子块赋值多维数组以提升性能

错误原因 #

你遇到的Cannot coerce list to type 'double complex'错误，是因为Cython内存视图（complex[:,:]）不支持直接将Python嵌套列表赋值给切片。内存视图的切片赋值仅接受兼容的numpy数组或其他内存视图，不能直接用Python原生列表。

极致性能实现方案 #

为了最大化性能，我们完全用C级别的操作完成赋值，避免任何Python层面的交互。以下是修改后的代码：

# test.pyx
import numpy as np
cimport numpy as cnp
cnp.import_array()

def HGamma_m( int L, double[:] param, cnp.ndarray[int64_t, ndim=2, mode='c'] R ):
    # 初始化内存视图，指定C连续布局
    cdef cnp.ndarray[cnp.complex128, ndim=2, mode='c'] HGamma_np = np.zeros( (4*L,4*L), dtype=np.complex128 )
    cdef complex[:,:] HGamma = HGamma_np
    
    cdef int i, j
    cdef double l, m, n
    cdef double bowler = param[0]  # 假设bowler是param的第一个元素，根据实际情况调整
    cdef double p4_p5 = param[4] - param[5]  # 预计算重复表达式
    cdef double val
    
    for i in range(L):
        for j in range(L):
            # 直接用内存视图访问R的元素，避免Python数组操作
            l = R[j, 0] - R[i, 0]
            m = R[j, 1] - R[i, 1]
            n = R[j, 2] - R[i, 2]
            
            # 逐个赋值4x4子块的每个元素
            # 第0行
            val = bowler * param[2]
            HGamma[4*i + 0, 4*j + 0] = val
            val = bowler * param[3] * l
            HGamma[4*i + 0, 4*j + 1] = val
            val = bowler * param[3] * m
            HGamma[4*i + 0, 4*j + 2] = val
            val = bowler * param[3] * n
            HGamma[4*i + 0, 4*j + 3] = val
            
            # 第1行
            val = -bowler * param[3] * l
            HGamma[4*i + 1, 4*j + 0] = val
            val = bowler * param[4] * l*l + bowler * param[5] * (1 - l*l)
            HGamma[4*i + 1, 4*j + 1] = val
            val = bowler * m * l * p4_p5
            HGamma[4*i + 1, 4*j + 2] = val
            val = bowler * n * l * p4_p5
            HGamma[4*i + 1, 4*j + 3] = val
            
            # 第2行
            val = -bowler * param[3] * m
            HGamma[4*i + 2, 4*j + 0] = val
            val = bowler * l * m * p4_p5
            HGamma[4*i + 2, 4*j + 1] = val
            val = bowler * param[4] * m*m + bowler * param[5] * (1 - m*m)
            HGamma[4*i + 2, 4*j + 2] = val
            val = bowler * n * m * p4_p5
            HGamma[4*i + 2, 4*j + 3] = val
            
            # 第3行
            val = -bowler * param[3] * n
            HGamma[4*i + 3, 4*j + 0] = val
            val = bowler * l * n * p4_p5
            HGamma[4*i + 3, 4*j + 1] = val
            val = bowler * m * n * p4_p5
            HGamma[4*i + 3, 4*j + 2] = val
            val = bowler * param[4] * n*n + bowler * param[5] * (1 - n*n)
            HGamma[4*i + 3, 4*j + 3] = val
    
    return HGamma_np  # 返回numpy数组方便后续使用

关键优化点 #

• 去掉Python列表赋值：改用C级别的逐个元素赋值，完全避免Python对象的创建和交互开销。
• 预计算重复表达式：比如param[4]-param[5]只计算一次，减少重复运算。
• 指定内存布局：给数组和内存视图加上mode='c'（C连续），确保内存访问是连续的，提升缓存命中率。
• 直接访问内存视图元素：R[j,0]代替R[j]-R[i]这种Python数组操作，避免临时数组的创建。
• 静态类型变量：所有变量都用cdef声明静态类型，消除Python动态类型的开销。

编译与验证 #

保持你的setup.py不变，编译后运行即可。这个版本的性能会远优于numpy实现，因为所有核心计算都在C层面完成，没有Python循环的开销。

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