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.

TDMAsolver.py

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

Test code:

A = np.array([[10,2,0,0],[3,10,4,0],[0,1,7,5],[0,0,3,4]],dtype=float)

a = np.array([3.,1,3])
b = np.array([10.,10.,7.,4.])
c = np.array([2.,4.,5.])
d = np.array([3,4,5,6.])

print("Test results:")
print(TDMAsolverold(a, b, c, d))
print(TDMAsolver(a, b, c, d))
print(np.linalg.solve(A, d))

print("Speed results:")
t0 = time.time()
for i in range(100000):
    TDMAsolver(a, b, c, d)
t1 = time.time()
print("jit_new {}".format(t1-t0))
t0 = time.time()
for i in range(100000):
    TDMAsolverold(a, b, c, d)
t1 = time.time()
print("old_without {}".format(t1-t0))
t0 = time.time()
for i in range(100000):
    np.linalg.solve(A, d)
t1 = time.time()
print("control {}".format(t1-t0))

Output on my laptop:

Test results:
[ 0.14877589  0.75612053 -1.00188324  2.25141243]
[ 0.14877589  0.75612053 -1.00188324  2.25141243]
[ 0.14877589  0.75612053 -1.00188324  2.25141243]
Speed results:
jit_new 0.47366952896118164
old_without 1.140716314315796
control 1.125091552734375

These results are promising as the speed is increased by a factor of 2 approximately. I'm solving the tridiagonal matrix approximately millions of times in my simulation so this is a huge improvement over what I had before.

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.

I don't understand why, but I got problems when trying to test it on small matrices:
(and jit is very slow for my tests...)

a = np.array([2, 2])
b = np.array([1, 1, 1])
c = np.array([ 2, 2])
d = np.array([ 1, 1, 1])
print(TDMAsolver(a, b, c, d))
A = np.array([[1, 2, 0], [2, 1, 2], [0, 2, 1]], dtype=float)
print(np.linalg.solve(A, d))

TDMA gives me [1 0 0], while linalg.solve fives me the good answer [ 0.14285714 0.42857143 0.14285714]
Any idea why TDMA is wrong here?

@dl-wuhee
It might be with jit. jit compiles the code to C but it works better with some processors. Best to read the site of numba on this.

@MarineLasbleis
You are right, something goes wrong here. It seems to start working from 4x4 matrix.

Thanks for this 👍

Can you add a license file?

@ahwillia The file is forked from two other works that did not have any license file. I think you need to go to the original work ask him to license it and then I can add a license based on his license. Currently, under Github policy I can fork code and make it my own thing but since it is a derivative of some original code I think that is way it needs to go. But I'm no lawyer so, if you can advise me on this would be awesome. If I can I would release it under the most open license I can find.

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.