优化嵌套坐标变换：基于Numpy的演化算法加速方案问询

我正在使用Numpy和Pandas实现一种演化算法（Evolutionary Algorithm），该代码基于MathWorks平台上的一个实现方案。代码可正常运行并输出合理结果，但运行速度极慢。

核心代码片段如下：

# Initialize
xm = np.zeros((n, nlambda))
sigmam, alpham = list(np.ones(nlambda) * np.nan), list(np.ones(nlambda) * np.nan)
sigmam: List[np.ndarray]  # list of NUMBER OF OFFSPRINGS np arrays 264x264
alpham: List[np.ndarray]  # list of NUMBER OF OFFSPRINGS np arrays 264x264

for i in population_size:
   ### Some code ###

   alpham[i] = np.where(alpham[i] > math.pi, alpham[i] - 2 * math.pi, alpham[i])
   alpham[i] = np.where(alpham[i] < -math.pi, alpham[i] + 2 * math.pi, alpham[i])
   R = np.eye(n)
   for m in range(n - 1):
      for q in range(m + 1, n):
         T = np.eye(n)
         T[np.ix_([m, q], [m, q])] = [[np.cos(alpham[i][m,m]), -np.sin(alpham[i][m,q])], [np.sin(alpham[i][q,m]), np.cos(alpham[i][q,q])]]
         R[:, [m, q]] = np.dot(R, T[:, [m,q]])
    ### Some code ###

现寻求通过更巧妙地运用Numpy数组/向量化操作来优化该代码的可行方法。

优化途径 #

1. 向量化角度归一化操作 #

原代码两次np.where的角度约束可以用np.mod一次性完成，避免多次数组遍历：

# 替换原有的两次np.where
alpham[i] = np.mod(alpham[i] + np.pi, 2 * np.pi) - np.pi

该操作直接对整个数组进行矢量化计算，比两次条件判断效率更高。

2. 重构旋转矩阵更新逻辑 #

嵌套循环生成旋转矩阵是性能瓶颈核心，可通过直接操作矩阵列替代生成全量单位矩阵，减少内存开销和计算量：

# 初始化R为单位矩阵
R = np.eye(n)
cos_vals = np.cos(alpham[i])
sin_vals = np.sin(alpham[i])

for m in range(n-1):
    for q in range(m+1, n):
        # 提取当前旋转所需的三角函数值
        c1, s1 = cos_vals[m, q], sin_vals[m, q]
        c2, s2 = sin_vals[q, m], cos_vals[q, q]
        
        # 直接更新R的m、q列，避免生成完整T矩阵
        col_m = R[:, m].copy()
        col_q = R[:, q].copy()
        R[:, m] = c1 * col_m - s1 * col_q
        R[:, q] = s2 * col_m + c2 * col_q

3. 用三维数组替代列表存储 #

将sigmam和alpham从列表转换为三维数组（shape: (nlambda, n, n)），利用Numpy批量矢量化操作替代Python循环：

# 初始化时替换列表为三维数组
sigmam = np.full((nlambda, n, n), np.nan)
alpham = np.full((nlambda, n, n), np.nan)

# 批量处理所有个体的角度归一化
alpham = np.mod(alpham + np.pi, 2 * np.pi) - np.pi

这种方式能利用Numpy内部C实现处理批量数据，消除Python循环的开销。

4. 用Numba加速核心循环 #

若嵌套循环无法完全矢量化，可通过Numba的JIT编译将Python代码转为机器码执行：

from numba import jit

@jit(nopython=True)
def compute_rotation(n, alphai):
    R = np.eye(n)
    for m in range(n-1):
        for q in range(m+1, n):
            c1 = np.cos(alphai[m, q])
            s1 = np.sin(alphai[m, q])
            c2 = np.cos(alphai[q, q])
            s2 = np.sin(alphai[q, m])
            col_m = R[:, m].copy()
            col_q = R[:, q].copy()
            R[:, m] = c1 * col_m - s1 * col_q
            R[:, q] = s2 * col_m + c2 * col_q
    return R

# 主循环中调用编译后的函数
R = compute_rotation(n, alpham[i])

nopython模式可消除Python解释器开销，对数值密集型循环能带来数倍至数十倍的速度提升。

5. 减少不必要的数组拷贝 #

避免生成临时数组，直接操作原数组切片。例如原代码中T[:, [m,q]]会产生临时数组，优化后直接更新R的对应列，减少内存分配和拷贝成本。

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