使用NumPy实现三对角矩阵算法(TDMA)

Question

我正在用Python和NumPy实现TDMA（时分多址）。我把三对角矩阵存储在三个数组里：

a = array([...])
b = array([...])
c = array([...])

我想高效地计算alpha系数。算法大致如下：

# n = size of the given matrix - 1
alpha = zeros(n)
alpha[0] = b[0] / c[0]
for i in range(n-1):
    alpha[i+1] = b[i] / (c[i] - a[i] * alpha[i])

不过，这种方法效率不高，因为Python的for循环比较慢。我想要的是类似这样的做法：

# n = size of the given matrix - 1
alpha = zeros(n)
alpha[0] = b[0] / c[0]
alpha[1:] = b[1:] / (c[1:] - a[1:] * alpha[:-1])

但是在这种情况下，结果是不正确的，因为NumPy会把最后一个表达式的右边部分存储在一个临时数组中，然后把这个临时数组的元素引用赋值给alpha[1:]。所以a[1:] * alpha[:-1]的结果就变成了一个全是零的数组。

有没有办法让NumPy在它的内部循环中使用之前步骤计算出来的alpha的值呢？

谢谢。

Answer 1

如果你想解决三对角系统，可以使用 solve_banded() 这个函数，它在 numpy.linalg 这个库里。具体是不是你需要的，我就不太确定了。

Answer 2

显然，要在Python中做到这一点，没有其他办法，只能使用C语言或者它的“Python风格”的变种。