用NumPy实现三对角矩阵算法（TDMA）

Question

我正在使用NumPy在Python中实现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将最后一个表达式的右侧部分存储在temprorary数组中，然后将对其元素的引用分配给alpha[1:]。因此a[1:] * alpha[:-1]只是一个零数组。

有没有办法告诉NumPy使用在其内部循环中先前步骤计算的alpha值？

感谢。

Answer 1

如果要解决的三对角系统在solve_banded()中有numpy.linalg。不确定这是不是你想要的。

Answer 2

显然，如果不使用C或其 pythonic 变体，就无法在Python中执行此操作。