如何用Cython进一步加速数组循环？光学计算性能优化咨询

我开发了一款光学工具，需要针对不同时间步长和波长计算一组复振幅（波前），核心逻辑是两层循环，每次循环调用下方的Cython方法。目前通过Cython仅实现了3倍提速，单次方法调用仍需0.4ms，无法满足需求。通常时间步数组长度为100，波长数组长度为30，总计需调用3000次。我曾考虑改写方法以接收数组批量传入时间和波长，但需大量使用numpy.einsum()，听说该函数难以通过Cython优化。现提出两个问题：

1. 能否通过其他方式用Cython进一步加速该方法？
2. 传入二维数组替代两层循环是否在计算上更合理？

当前Cython代码 #

cimport numpy as cnp
cimport cython
import numpy as np
from libc.math cimport exp, cos, sin

ctypedef cnp.complex128_t DTYPE_t

@cython.cdivision(True)
@cython.boundscheck(False) # compiler directive
@cython.wraparound(False) # compiler directive
cpdef DTYPE_t[:,:,::1] _get_input_complex_amplitudes(
                                  float time,
                                  float wavelength,
                                  double[::1] x_source_sky_coordinates,
                                  double[::1] y_source_sky_coordinates,
                                  double[::1] x_observatory_coordinates,
                                  double[::1] y_observatory_coordinates,
                                  float aperture_radius,
                                  int number_of_inputs,
                                  int grid_size,
                                  DTYPE_t[:,:,::1] out):
    cdef unsigned int ix, iy, index_input
    cdef double arg, cos_val, sin_val
    cdef double factor = 2 * 3.1415926536
    cdef DTYPE_t[:,:,::1] ou = out

    for index_input in range(number_of_inputs):
        for ix in range(grid_size):
            for iy in range(grid_size):
                arg = factor / wavelength * (
                        x_observatory_coordinates[index_input] * x_source_sky_coordinates[ix] +
                        y_observatory_coordinates[index_input] * y_source_sky_coordinates[iy])

                cos_val = cos(arg)
                sin_val = sin(arg)

                ou[index_input,ix, iy] = aperture_radius * (cos_val + 1j * sin_val)
    return ou

问题1：Cython进一步加速的方案 #

可以通过以下几个方向针对性优化：

• 替换三角函数组合为复指数计算：原代码分开计算cos和sin再组合复数，直接用C99的复数指数函数cexp更高效。从libc.math cimport cexp，然后将赋值语句改为ou[index_input,ix, iy] = aperture_radius * cexp(1j * arg)，减少一次函数调用开销。
• 提前计算常量值：factor / wavelength在循环内重复计算，提前在函数开头算出inv_wavelength = factor / wavelength，循环内直接使用该变量，减少除法运算次数。同时用libc.math中的M_PI替代硬编码的圆周率，既精确又可能提升编译优化空间。
• 减少数组索引访问：在index_input循环内，提前把x_observatory_coordinates[index_input]和y_observatory_coordinates[index_input]赋值给局部变量，避免循环内重复索引数组：

  cdef double x_obs = x_observatory_coordinates[index_input]
  cdef double y_obs = y_observatory_coordinates[index_input]
  

• 开启编译器高级优化：在编译Cython代码时，添加-O3和-ffast-math编译参数（如在setup.py的extra_compile_args中配置），让编译器做更激进的数学运算优化。
• 清理冗余参数：当前代码中time参数未被使用，直接删除以减少函数调用的参数传递开销。

问题2：批量传入数组是否更合理 #

是的，批量处理在计算效率上更优，核心原因如下：

• 消除Python-C调用开销：3000次函数调用会累积大量Python与C的切换开销，批量处理只需一次调用，能节省这部分可观的时间。
• 无需依赖numpy.einsum：可以直接在Cython中扩展循环逻辑，接收float[:] wavelengths和float[:] times数组，新增循环遍历所有波长和时间步，直接在C层完成批量计算；也可以用NumPy的广播特性实现向量化计算，例如：

  # 批量计算所有波长的示例
  arg = 2 * np.pi / wavelengths[:, None, None, None] * (
      x_observatory_coordinates[None, :, None, None] * x_source_sky_coordinates[None, None, :, None] +
      y_observatory_coordinates[None, :, None, None] * y_source_sky_coordinates[None, None, None, :]
  )
  out = aperture_radius * np.exp(1j * arg)
  

  这种向量化操作基于NumPy的底层C实现，性能大概率优于循环调用Cython函数。
• 内存复用更高效：批量处理时可以预先分配大尺寸的输出数组，避免多次内存分配与释放的开销，进一步提升整体效率。

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