Skip to content

Instantly share code, notes, and snippets.

@TheoChristiaanse
Forked from cbellei/TDMAsolver.py
Last active November 15, 2021 14:19
Show Gist options
  • Save TheoChristiaanse/d168b7e57dd30342a81aa1dc4eb3e469 to your computer and use it in GitHub Desktop.
Save TheoChristiaanse/d168b7e57dd30342a81aa1dc4eb3e469 to your computer and use it in GitHub Desktop.
Tridiagonal Matrix Algorithm solver in Python. I've modified the code from cbellei so, it works with python 3.0+ and implemented the use of jit to increase the speed.
import numpy as np
from numba import jit, f8
## Tri Diagonal Matrix Algorithm(a.k.a Thomas algorithm) solver
@jit(f8[:] (f8[:],f8[:],f8[:],f8[:] ))
def TDMAsolver(a, b, c, d):
'''
TDMA solver, a b c d can be NumPy array type or Python list type.
refer to http://en.wikipedia.org/wiki/Tridiagonal_matrix_algorithm
and to http://www.cfd-online.com/Wiki/Tridiagonal_matrix_algorithm_-_TDMA_(Thomas_algorithm)
'''
nf = len(d) # number of equations
ac, bc, cc, dc = map(np.array, (a, b, c, d)) # copy arrays
for it in range(1, nf):
mc = ac[it-1]/bc[it-1]
bc[it] = bc[it] - mc*cc[it-1]
dc[it] = dc[it] - mc*dc[it-1]
xc = bc
xc[-1] = dc[-1]/bc[-1]
for il in range(nf-2, -1, -1):
xc[il] = (dc[il]-cc[il]*xc[il+1])/bc[il]
return xc
@richinex
Copy link

I run your code, and got different speed results.

Test results:
[ 0.14877589 0.75612053 -1.00188324 2.25141243]
[ 0.14877589 0.75612053 -1.00188324 2.25141243]
Speed results:
jit_new 6.220152854919434
control 0.7551865577697754

I am not familiar with the numba or jit. So I don't know if it is caused by this. Anyway, thank you for your improvement.

Indeed the code gives wrong answers. I tried to compare the results also.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment