Tridiagonal Matrix Algorithm solver in Python, using Numpy array - http://ofan666.blogspot.com/2012/02/tridiagonal-matrix-algorithm-solver-in.html

TDMAsolver.py

try:
   import numpypy as np    # for compatibility with numpy in pypy
except:
   import numpy as np      # if using numpy in cpython

## 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
    '''
    nf = len(a)     # number of equations
    ac, bc, cc, dc = map(np.array, (a, b, c, d))     # copy the array
    for it in xrange(1, nf):
        mc = ac[it]/bc[it-1]
        bc[it] = bc[it] - mc*cc[it-1] 
        dc[it] = dc[it] - mc*dc[it-1]

    xc = ac
    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]

    del bc, cc, dc  # delete variables from memory

    return xc

How does this compare to a C version? I wrote a C extension to Python of this algorithm that inverts a 100,000,000 element tridiagonal matrix in about 3 milliseconds. Check out my Github page for more details.

@cpcloud THX, of course your C-version will lot faster, in my machine it fast it is 1.691 ms:
here your's PyTDMA in my machine:

result test PyTDMA

~/cpcloud-PyTDMA-32772f4$ ./test.py 10000
error: 4.762e-12, runtime: 1.691 ms, 100,000,000 elements

and here is mine using numpy linked w/ Intel MKL in cpython

using IPython %timeit

%timeit TDMAsolver(a, b, c, d)
10 loops, best of 3: 30.9 ms per loop

using time

$ time python TDMApypy.py

real 0m0.129s
user 0m0.100s
sys 0m0.020s

and this is when using numpypy in pypy

$ time pypy TDMApypy.py

real 0m0.093s
user 0m0.080s
sys 0m0.010s

After creating the test matrix and vectors, if I run %timeit tdma(A, rhs) the result is 1000 loops, best of 3: 788 us per loop. A and rhs are 10000 x 10000 and 10000 elements, respectively.

I'm also not entirely sure which is the best way to run it. In the test.py script I'm using the Timer class from the timeit module, and passing an anonymous function.

According to the Python documentation, passing lambdas to the Timer class has more overhead than providing a string, but you can't really provide a string of a large NumPy array. Any thoughts?

The %timeit magic function in IPython gives nice results :) but I want to be sure that it's measuring the computation speed accurately. Thanks.

I'm thinking that calling the UNIX time function is not the way to go because it times everything, which is not what I'm interested in: I just want to know the speed of the solution, not how long it takes to create a huge matrix and bunch of vectors in addition to the solution. There's a ton of overhead in just creating a big matrix.

The reason you're getting a memory error (I presume) is because 100,000 * 100,000 * 8 is the number of bytes you need to store a square matrix of size 100,000 by 100,000. So, you'd need 80GB of RAM (!) to store that beast in memory all at once. That's not including RAM you'd need for anything else.

In general, for an array with 64-bit floating point numbers the amount of RAM needed to store a power of ten more elements increases cubicly.

Hello Everyone
Would you please help me with this question
If I have this matrix
[a2 1 0 ] [x2^n+1 ] = X2^n - X1
[1 a3 1 ] [x3^n+1 ] = X3^n
[0 1 a4] [x4^n+1 ] = X4^n - X5

Where X1 and X5 are given
How can I solve it with Thomas Algorithm using Python
Thanx in advance
Best Regards

@ofan666
Thanks for your code. I have corrected a couple of bugs, see here
https://gist.github.com/cbellei/8ab3ab8551b8dfc8b081c518ccd9ada9

Could you add a license file to this code?

Could you add a license file to this code?

I agree. I appreciate the code snippet, but I would be happier if there was an explicit license. (I'm using this short code too.) Fortunately, it's one of the easiest methods in linear algebra, so it is easy to rewrite.