Cython中Jacobi迭代函数返回多个对象的最高效实现方案

问题背景 #

我正在开发一款2D拉普拉斯方程（Laplace.eq）求解器，在收敛检查环节，函数需要返回多个对象：处理后的数组、迭代次数以及当前估计误差。以下是Jacobi方法的最简实现代码：

@cython.boundscheck(False)
@cython.wraparound(False)
cdef void jacobi_iteration(double[:,:] A, double[:,:] mask):
    cdef int i, j, n, m
    cdef double tmp
    n, m = A.shape[0], A.shape[1]
    for i in prange(1, n-1, nogil=True):
        for j in range(1, m-1):
            if not mask[i, j]:
                tmp = 0.25 * (A[i-1, j] + A[i+1, j] + A[i, j-1] + A[i, j+1])
                A[i, j] = tmp

@cython.boundscheck(False)
@cython.wraparound(False)
cpdef Jacobi_fx(double[:,:] A, double[:,:] mask, int max_iterations, double error):
    cdef int i, n, m, iteration
    cdef double norm, norm_prev, tmp_error
    norm_prev = 0
    norm = 0
    iteration = 0
    n, m = A.shape[0], A.shape[1]
    for i in range(max_iterations):
        norm = np.linalg.norm(A)
        jacobi_iteration(A, mask)
        tmp_error = abs(norm - norm_prev)
        norm_prev = norm
        iteration = iteration + 1
        if tmp_error < error:
            break
    return A, tmp_error, iteration

其中mask数组用于Dirichlet边界条件，标记需要迭代的索引。当前函数返回A、tmp_error、iteration的方式在注解中显示解释开销最大。我考虑过以下几种替代方案：

• 返回列表：

  cpdef list Jacobi_fx(...)
  

  但这种方式通过append添加对象速度较慢；
• 返回多维数组：

  cpdef double[:,::1][3] Jacobi_fx(...)
  

  不确定语法是否正确，意图是返回包含3个二维数组的一维数组，虽没必要但可能比append快；
• 使用C struct：

  cdef struct Data:
      double[:, :] A
      double tmp_error
      int iteration
  

  对应函数实现：

  cpdef Data Jacobi_fx(...):
      ...#循环前代码不变
      data = Data(A=A, tmp_error=0.0, iteration=0)
  
      for i in range(max_iterations):
          norm = np.linalg.norm(data.A)
          jacobi_iteration(data.A, mask)
          data.tmp_error = abs(norm - norm_prev)
          norm_prev = norm
          data.iteration = iteration + 1
          if data.tmp_error < error:
              break
      return data
  

• 使用类实现，但因脚本较短想避免使用类。

想请教：在Cython中返回多个对象的最高效方法是什么？尤其是针对会被多次调用的函数。

解答 #

针对你的场景，使用C struct是最高效的方案，理由如下：

1. 最小化Python交互开销：C struct是纯C层面的结构，返回时不需要像元组/列表那样进行Python对象的包装与转换，完全避免了跨C-Python边界的解释开销——这正是你当前返回元组时遇到的主要性能瓶颈。
2. 语法与实现简洁：你给出的struct定义和函数实现是正确的，不需要额外的复杂逻辑，直接将数据打包到struct中返回即可，代码可读性和维护性都不错。
3. 性能碾压其他方案：
   • 返回列表需要多次调用Python的append方法，每次都涉及Python API调用，性能损耗明显；
   • 返回多维数组的语法并不正确（Cython不支持直接返回这种嵌套的数组视图），而且强行将标量（误差、迭代次数）塞进数组会造成内存浪费和不必要的类型转换；
   • 类的开销比struct大得多，类实例涉及Python对象的创建、属性访问的额外逻辑，完全不符合你追求高效的需求。

另外，还有一个优化点：你的当前代码中np.linalg.norm(A)会触发Python调用，如果你追求极致性能，可以自己实现一个C层面的范数计算函数，避免跨边界的开销。

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