三對角矩陣算法（TDMA）又名托馬斯算法，在NumPy數組中使用Python

Question

我在MATLAB中找到了thomas算法或TDMA的實現。

function x = TDMAsolver(a,b,c,d) 
    %a, b, c are the column vectors for the compressed tridiagonal matrix, d is the right vector 
    n = length(b); % n is the number of rows 

    % Modify the first-row coefficients 
    c(1) = c(1)/b(1); % Division by zero risk. 
    d(1) = d(1)/b(1); % Division by zero would imply a singular matrix. 

    for i = 2:n-1 
     temp = b(i) - a(i) * c(i-1); 
     c(i) = c(i)/temp; 
     d(i) = (d(i) - a(i) * d(i-1))/temp; 
    end 

    d(n) = (d(n) - a(n) * d(n-1))/(b(n) - a(n) * c(n-1)); 

    % Now back substitute. 
    x(n) = d(n); 
    for i = n-1:-1:1 
     x(i) = d(i) - c(i) * x(i + 1); 
    end 
end 

我需要它在python中使用numpy數組，這裏我第一次嘗試在python中的算法。

import numpy 

aa = (0.,8.,9.,3.,4.) 
bb = (4.,5.,9.,4.,7.) 
cc = (9.,4.,5.,7.,0.) 
dd = (8.,4.,5.,9.,6.) 

ary = numpy.array 

a = ary(aa) 
b = ary(bb) 
c = ary(cc) 
d = ary(dd) 

n = len(b)## n is the number of rows 

## Modify the first-row coefficients 
c[0] = c[0]/ b[0] ## risk of Division by zero. 
d[0] = d[0]/ b[0] 

for i in range(1,n,1): 
    temp = b[i] - a[i] * c[i-1] 
    c[i] = c[i]/temp 
    d[i] = (d[i] - a[i] * d[i-1])/temp 

d[-1] = (d[-1] - a[-1] * d[-2])/(b[-1] - a[-1] * c[-2]) 

## Now back substitute. 
x = numpy.zeros(5) 
x[-1] = d[-1] 
for i in range(-2, -n-1, -1): 
    x[i] = d[i] - c[i] * x[i + 1] 

他們給出了不同的結果，所以我做錯了什麼？

Answer 1

有在兩者之間的至少一個差異：

for i in range(1,n,1): 

在Python迭代從索引1至最後一個索引 n-1

，而從索引1

for i = 2:n-1 

迭代（從零開始）到（last-1）索引，因爲Matlab有一個基於索引的索引。

Answer 2

在你的循環中，Matlab版本遍歷第二個到第二個到最後一個元素。做同樣在Python，你想：

for i in range(1,n-1): 

（如voithos的評論指出，這是因爲範圍功能不包括最後一個索引，所以你需要除了糾正這一點的變化，以0索引）。

Answer 3

我做了這個，因爲沒有任何python在線實現實際工作。我已經測試了它與內置矩陣求逆和結果匹配。

這裏，=更低Diag（診斷），B =主Diag（診斷），C =上診斷，d =解向量

import numpy as np 

def TDMA(a,b,c,d): 
    n = len(d) 
    w= np.zeros(n-1,float) 
    g= np.zeros(n, float) 
    p = np.zeros(n,float) 

    w[0] = c[0]/b[0] 
    g[0] = d[0]/b[0] 

    for i in range(1,n-1): 
     w[i] = c[i]/(b[i] - a[i-1]*w[i-1]) 
    for i in range(1,n): 
     g[i] = (d[i] - a[i-1]*g[i-1])/(b[i] - a[i-1]*w[i-1]) 
    p[n-1] = g[n-1] 
    for i in range(n-1,0,-1): 
     p[i-1] = g[i-1] - w[i-1]*p[i] 
    return p 

對於大型基質，使用numba一個容易的性能提升！這段代碼在我的測試中勝過np.linalg.inv（）：

import numpy as np 
from numba import jit  

@jit 
def TDMA(a,b,c,d): 
    n = len(d) 
    w= np.zeros(n-1,float) 
    g= np.zeros(n, float) 
    p = np.zeros(n,float) 

    w[0] = c[0]/b[0] 
    g[0] = d[0]/b[0] 

    for i in range(1,n-1): 
     w[i] = c[i]/(b[i] - a[i-1]*w[i-1]) 
    for i in range(1,n): 
     g[i] = (d[i] - a[i-1]*g[i-1])/(b[i] - a[i-1]*w[i-1]) 
    p[n-1] = g[n-1] 
    for i in range(n-1,0,-1): 
     p[i-1] = g[i-1] - w[i-1]*p[i] 
    return p