如何用向量化操作替换填充矩阵时的嵌套for循环以优化性能

开发者阿泽

2026-6-29

用向量化操作替代嵌套循环优化矩阵填充性能

原代码中针对i和j的嵌套循环会带来显著性能开销，尤其是当N_Om数值较大时，Python级别的循环会大幅拖慢执行速度。可以利用NumPy的向量化索引与广播特性，完全替代这两层循环，实现高效的矩阵填充。

优化后的代码如下（替换原有的i、j嵌套循环部分）：

import numpy as np

dim = N_Om * 2 * a
A = np.zeros((dim, dim), dtype="complex")  # 简化初始化写法

for alpha in range(0, 2):
    for l in range(1, l_max + 1):
        B = get_B(l, alpha)  # small matrix of rank NOm
        block_size = 2 * l + 1
        base_offset = alpha * a + (l - 1) * (l + 1)

        # 生成向量化索引数组
        i_arr = np.arange(N_Om)
        k_arr = np.arange(block_size)

        # 计算每个块的行/列起始位置
        row_starts = i_arr * 2 * a + base_offset
        col_starts = i_arr * 2 * a + base_offset

        # 构造三维索引：定位所有需要赋值的对角线元素位置
        rows = row_starts[:, None, None] + k_arr[None, None, :]
        cols = col_starts[None, :, None] + k_arr[None, None, :]

        # 批量赋值，替代嵌套循环
        A[rows, cols] = B[:, :, None]

代码逻辑说明

先计算当前l对应的对角块大小block_size，以及由alpha和l决定的基础偏移量base_offset。
用np.arange生成i（块索引）和k（块内对角位置）的索引数组，避免Python层面的循环。
通过广播构造行、列索引：rows形状为(N_Om, 1, block_size)，cols形状为(1, N_Om, block_size)，两者广播后能精确定位到A中所有需要填充的对角线元素位置。
最后通过A[rows, cols] = B[:, :, None]完成批量赋值，将B[i,j]的值填充到对应块的所有对角线上，完全等价于原循环中B[i,j] * np.eye(block_size)的操作。