Tridiagonal Matrix Algorithm solver in Python

TDMAsolver.py

import numpy as np

## Tri Diagonal Matrix Algorithm(a.k.a Thomas algorithm) solver
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 xrange(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 xrange(nf-2, -1, -1):
        xc[il] = (dc[il]-cc[il]*xc[il+1])/bc[il]

    return xc

Example:

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 TDMAsolver(a, b, c, d)
>> [ 0.14877589  0.75612053 -1.00188324  2.25141243]
#compare against numpy linear algebra library
print np.linalg.solve(A, d)
>> [ 0.14877589  0.75612053 -1.00188324  2.25141243]

Thanks for this! I've forked your gist to make it work in 3.0 and increase the speed by a factor of 2 by using numba.

License? (I also nagged ofan666)
(Thanks for posting this, by the way.)

I've added a GNU General Public License as a separate gist. Does this help? (thanks for checking by the way!)

Thanks! That helps.

How to be in case when on diagonal line during computing we get zero?
for example:
A = np.array([[6,2,0,0],[3,1,4,0],[0,1,7,5],[0,0,3,4]],dtype=float)

How to be in case when on diagonal line during computing we get zero?
for example:
A = np.array([[6,2,0,0],[3,1,4,0],[0,1,7,5],[0,0,3,4]],dtype=float)

@BatoNam, there is condition, when you could use this method. Conditional is
abs(a[i]) + abs(c[i]) <= abs(b[i]) for each i

xc is just a reference to bc
It's pointless.