Tridiagonal Matrix Algorithm (TDMA) aka Thomas Algorithm, using Python with NumPy arrays
23,445
Solution 1
There's at least one difference between the two:
for i in range(1,n,1):
in Python iterates from index 1 to the last index n-1
, while
for i = 2:n-1
iterates from index 1
(zero-based) to the last-1
index, since Matlab has one-based indexing.
Solution 2
I made this since none of the online implementations for python actually work. I've tested it against built-in matrix inversion and the results match.
Here a = Lower Diag
, b = Main Diag
, c = Upper Diag
, d = solution vector
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
For an easy performance boost for large matrices, use numba! This code outperforms np.linalg.inv()
in my tests:
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
Author by
Admin
Updated on January 05, 2021Comments
-
Admin over 3 years
I found an implementation of the thomas algorithm or TDMA in MATLAB.
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
I need it in python using numpy arrays, here my first attempt at the algorithm in 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]
They give different results, so what am I doing wrong?