Numba @njit加速Thomas算法三对角求解器报非恒定值捕获错误

• Numba的nopython模式不支持逃逸闭包捕获外部非常量变量的场景：你在tridiag函数内部定义的a、eff、u三个嵌套函数，都捕获了x1、y、z、alpha等运行时才确定值的非常量变量，且这三个函数被作为参数传递给了accumulate（即函数发生逃逸），直接触发了本次报错。
• itertools.accumulate搭配自定义函数的用法在Numba nopython模式下兼容性极差，本身也不推荐在@njit装饰的函数中使用。

直接手写Thomas算法的显式循环实现，完全规避闭包和accumulate的使用，Numba对原生循环的优化效率很高，写法也更易维护。

import numpy as np
from numba import njit

@njit
def tridiag(x, y, z, b):
    n = len(b)
    # 初始化结果数组
    u = np.zeros(n, dtype=np.float64)
    # 预处理x1
    x1 = np.concatenate((np.array([0.0], dtype=np.float64), np.array(x, dtype=np.float64)))
    # 前向扫描计算alpha和f数组，替代原先的accumulate逻辑
    alpha = np.zeros(n, dtype=np.float64)
    f = np.zeros(n, dtype=np.float64)
    alpha[0] = y[0]
    f[0] = b[0] / y[0]
    for j in range(1, n):
        alpha[j] = y[j] - x1[j] * z[j-1] / alpha[j-1]
        f[j] = (b[j] - x1[j] * f[j-1]) / alpha[j]
    # 反向回代
    u[-1] = f[-1]
    for j in range(n-2, -1, -1):
        u[j] = f[j] - z[j] * u[j+1] / alpha[j]
    return u

# 测试代码
x = np.array([-1 for i in range(9)], dtype=np.float64)
y = np.array([2 for i in range(10)], dtype=np.float64)
z = np.array([-1 for i in range(9)], dtype=np.float64)
b = np.array([1,1,1,1,1,1,1,1,1,1.5], dtype=np.float64)

print(tridiag(x,y,z,b))
# 输出结果与预期一致：
# [ 5.04545455  9.09090909 12.13636364 14.18181818 15.22727273
#  15.27272727 14.31818182 12.36363636  9.40909091  5.45454545]

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