含零对角块的分块对角矩阵高效求逆方法问询

场景说明：
现有分块对角矩阵（部分对角块为零），为某算子的表示形式，需求解其逆矩阵（非零块条件数足够低，无数值稳定性顾虑）。实际应用中N约为1000，且零块位置已知，希望逆矩阵的零块保持为零。当前采用逐个求逆非零块再填充的方法效率较低，寻求更高效实现。

核心优化思路 #

避免Python层面的循环开销，利用numpy的批量向量化操作，将非零块的提取、求逆、赋值全部转为底层C实现的操作，大幅提升效率。

实现代码 #

import numpy as np

np.random.seed(42)
N = 1000  # 实际场景规模
block_size = 3  # 对角块的固定大小

# 生成模拟的分块对角矩阵（零块位置已知：第0个对角块为零）
def composite_operator(N, block_size):
    matrix = np.zeros((block_size * N, block_size * N))
    for n in range(1, N):
        start_idx = block_size * n
        matrix[start_idx:start_idx+block_size, start_idx:start_idx+block_size] = np.random.rand(block_size, block_size)
    return matrix

# 生成原矩阵
A = composite_operator(N, block_size)

# 高效求解逆矩阵
inv_A = np.zeros_like(A)

# 1. 生成所有非零块的元素索引（利用已知的零块位置）
non_zero_block_starts = np.arange(1, N) * block_size  # 非零块的起始行/列索引
# 展开为所有非零块的元素行/列索引
rows = (non_zero_block_starts[:, None] + np.arange(block_size)).flatten()
cols = rows.copy()

# 2. 批量提取所有非零块为三维数组：形状为(M, block_size, block_size)，M为非零块数量
non_zero_blocks = A[rows][:, cols].reshape(-1, block_size, block_size)

# 3. 批量求逆：numpy的linalg.inv支持三维数组，自动对每个二维子块求逆
inv_non_zero_blocks = np.linalg.inv(non_zero_blocks)

# 4. 批量将逆块元素赋值回结果矩阵
inv_A[rows, cols] = inv_non_zero_blocks.flatten()

效率提升原因 #

• 原方法的Python循环会带来大量解释器层面的开销（1000次循环的调用、切片操作），而批量操作将所有逻辑转为numpy底层优化的C代码执行，减少了Python与C之间的交互开销。
• np.linalg.inv对三维数组的批量处理内部做了并行化优化，比逐个求逆的累计速度更快。

额外说明 #

如果零块的位置不是连续的（比如分散在不同位置），只需调整non_zero_block_starts的生成逻辑即可，只要已知非零块的起始索引，就能复用这套批量处理逻辑。

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