如何提升numpy线性求解循环运算的计算效率

优化方案 #

你当前的性能瓶颈主要来自Python层面的循环开销，以及每次循环中重复构造小数组、重复调用np.linalg.solve的额外成本，可以通过向量化运算完全消除循环，大幅提升速度：

方案1：使用numpy原生批量求解（通用适配任意小矩阵尺寸）

np.linalg.solve本身支持批量输入：当系数矩阵A为(N, m, m)的3D数组、右端项B为(N, m)的2D数组时，会自动对N个m×m的方程组并行求解，底层完全由C实现，没有Python循环开销。
代码示例：

import numpy as np

# 1. 批量构造3D系数矩阵A，shape为(2000, 2, 2)
A = np.stack([
    np.stack([A_11, A_12], axis=-1),
    np.stack([A_21, A_22], axis=-1)
], axis=1)

# 2. 批量构造右端项B，shape为(2000, 2)
B = np.stack([B_1, B_2], axis=-1)

# 3. 一次完成所有方程组求解，转置后得到shape为(2, 2000)的X
X = np.linalg.solve(A, B).T

方案2：2×2矩阵专属解析解（速度最快）

如果你的场景固定是2×2的方程组，可以直接用解析解实现逐元素运算，性能比通用的批量求解还要高3~5倍：

# 逐元素计算矩阵行列式
det = A_11 * A_22 - A_12 * A_21
# 逐元素计算两个未知数的解
x1 = (A_22 * B_1 - A_12 * B_2) / det
x2 = (-A_21 * B_1 + A_11 * B_2) / det
# 直接堆叠得到结果X
X = np.stack([x1, x2], axis=0)

注意事项 #

两种方案的行为和你原循环版本完全一致，如果存在行列式为0的奇异矩阵，都会抛出线性代数错误，若需要处理奇异场景，可以替换为最小二乘求解函数np.linalg.lstsq的批量版本。

内容的提问来源于stack exchange，提问作者Timo-m