Thomas algorithm

This is a function that solves a triangular matrix using Thomas Algorithm. This function will be very useful in the future. I'm planing to use it to solve partial differential equations later on.

The Code:

Save to file in a file named: ChemE.py

def thomas(a,b,c,d):
 '''Uses Thomas algorithm for solving a tridiagonal matrix for n unknowns.
  a, b, and  are a list of the matrix entries
  Matrix form of:
  [b1 c1    ] [x1] [d1]
  |a2 b2 c2   | |x2|  |d2|
  | a3 b3 c3  | |x3|  |d3|
  |     | |' |= | '|
  |     | |' |  | '|
  [   an bn cn] |xn] [dn]
 '''
 # Finding the size of the matrix and determine n
 n = len(b)
 #print n # Used for debugging
 # Test the size of a and c
 if len(a) != n-1:
  print 'Wrong index size for a.\n A should have an index of', n-1, '\n Your a has ', len(a)
  exit()
 if len(c) != n-1:
  print 'Wrong index size for c.\n C should have an index of', n-1, '\n Your c has', len(c)
  exit()
 
 # Converting to float and appending 0.0 to c
 for i in range(0,len(a)):
  a[i] = float(a[i])
 for i in range(0,len(b)):
  b[i] = float(b[i])
 for i in range(0,len(c)):
  c[i] = float(c[i])
 for i in range(0,len(d)):
  d[i] = float(d[i])
 c.append(0.0) # Hack to make the function to work

 # Calculate p and q
 p = []; q= [] 
 p.append(c[0]/b[0]); q.append(d[0]/b[0])
 for j in range(1,n):
  pj = c[j]/(b[j] - a[j-1]* p[j-1])
  qj = (d[j] - a[j-1]*q[j-1])/(b[j] - a[j-1]* p[j-1])
  p.append(pj); q.append(qj)
 #print p,q # Used for debugging the code!

 # Back sub
 x = []; x.append(q[n-1])
 for j in range(n-2,-1,-1):
  xj = q[j] - p[j]*x[0]  # Value holder
  x.insert(0,xj)   # Building the list backwards

 # Return the value
 return x

Example code execution:

This example is take from SEPARATION PROCESS PRINCIPLES: Chemical and Biochemical Operations 3ed by Seader, Henley, and Roper. Problem number is 10.5 on page 406.

"Use the Thomas algorithm to solve the following tridiagonal-matrix equation for the x vector."
| -6 3 0  0 0  0 0 0| |x1|    |   0   |
|  3 -4.5 3 0  0 0 0| |x2|    |   0   |
|  0  1.5 -7.5 3 0 0| |x3| = | 100 |
|  0  0  4.5 -7.5   3| |x4|    |   0   |
|  0   0  0   4.5-4.5| |x5|    |   0   |

#!/usr/bin/python

# Import the function:
from ChemE import thomas
# Matrix
a= [3,1.5,4.5,4.5]
b= [-6,-4.5,-7.5,-7.5,-4.5] 
c= [3,3,3,3]
d= [0,0,100,0,0]
# Execute the function
x = thomas(a,b,c,d)
# Print the result
print x