Solving a tridiagonal system of linear equations using the Thomas algorithm

Question

I'm trying to write a function that can solve a tridiagonal system of linear equations using the Thomas algorithm. It basically solves the following equation. (Details can be found at the Wiki page here Tridiagonal matrix algorithm.)

$$ \begin{bmatrix} {b_ 1} & {c_ 1} & { } & { } & { 0 } \\ {a_ 2} & {b_ 2} & {c_ 2} & { } & { } \\ { } & {a_ 3} & {b_ 3} & \ddots & { } \\ { } & { } & \ddots & \ddots & {c_{n-1}}\\ { 0 } & { } & { } & {a_n} & {b_n}\\ \end{bmatrix} \cdot \begin{bmatrix} {x_ 1 } \\ {x_ 2 } \\ {x_ 3 } \\ \vdots \\ {x_n } \\ \end{bmatrix} = \begin{bmatrix} {d_ 1 } \\ {d_ 2 } \\ {d_ 3 } \\ \vdots \\ {d_n } \\ \end{bmatrix} $$

This can be done easily in various ways in Mathematica using the built-in functions such as Solve, LinearSolve, LUDecomposition, etc. Instead of using Mathematica's wonderful black box solvers, I have decided to follow the procedures in the algorithm and write my own version which I can better control.

Here is the procedure in detail:

$$ c'_i= \begin{cases} \begin{array}{lcl} \cfrac{c_i}{b_i} &&; i = 1 \\ \cfrac{c_i}{b_i - c'_{i - 1} a_i} &&; i = 2, 3, \ \dots, n-1 \\ \end{array} \end{cases} $$

$$ d'_i= \begin{cases} \begin{array}{lcl} \cfrac{d_i}{b_i} &&; i = 1 \\ \cfrac{d_i - d'_{i - 1} a_i}{b_i - c'_{i - 1} a_i} &&; i = 2, 3, \dots, n. \\ \end{array} \end{cases} $$

$$ \begin{array}{lcl} x_n&=&d'_n\\ x_i&=&d'_i-c'_ix_{i+1} \end{array} $$

I tried to write the function in Mathematica like this:

tridag[{a_, b_, c_}, d_] := 
 Module[{n = Length[d], c1 = Range[Length[d]], d1 = Range[Length[d]], x = Range[Length[d]]},
   c1[[1]] = c[[1]]/b[[1]];
   d1[[1]] = c1[[1]]/c[[1]]*d[[1]];
   Do[c1[[i]] = c[[i]]/(b[[i]] - a[[i]]*c1[[i - 1]]);
     d1[[i]] = c1[[i]]/c[[i]]*(d[[i]] - a[[i]] d1[[i - 1]]);, {i, 2, n - 1}];
   x[[n]] = d1[[n]];
   Do[x[[i]] = d1[[i]] - c1[[i]]*x[[i + 1]], {i, n - 1, 1, -1}];
 x]

where a, b, c are the three diagonal rows in the matrix.

For a test case of

A = {{1, 2, 0, 0, 0}, {2, 2, 3, 0, 0}, {0, 3, 3, 4, 0}, {0, 0, 4, 4, 5}, {0, 0, 0, 5, 5}};
R = {5, 15, 31, 53, 45};

my function gives the same answer as Mathematica's linear solver:

LinearSolve[A, R]

{1, 2, 3, 4, 5}

tridag[{PadLeft[Diagonal[A, 1], 5], Diagonal[A], PadRight[Diagonal[A, -1], 5]}, R]

{1, 2, 3, 4, 5}

Here are my questions:

1. Since I am just a novice, my function is simply a direct translation of the above procedurual equations. So I would like to learn a really functional Mathematica way of implementing this function (other than using the built-in solver)?
2. A more general question is that, given an algorithm written in a procedure way, i.e., with implementation in a procedure language like C or Fortran in mind, how can it be translated into functional Mathematica code? Can you give some examples or suggestions or guidelines?

Edit

As belisarius points out, "'the real elegant Mathematica way' is always to use an already implemented function". To be clear, I'm not asking how to implement a solver that is superior to the built-in solver. I'm just asking how to implement a solver in the Mathematica way, i.e., in the functional way. I think knowing how to write one's own solver is helpful when using the built-in solvers. It gives the user confidence when the built-in solver does not work. An example may be the HilbertMatrix:

LinearSolve[ HilbertMatrix [10], Table[Random[], {10}]]

LinearSolve::luc: "Result for LinearSolve of badly conditioned matrix {{1., 0.5, 0.333333, 0.25, 0.2, 0.166667, 0.142857, 0.125, 0.111111, 0.1}, <<8>>, {0. 1, 0.0909091, 0.0833333, 0. 0769231, 0.0714286, 0.0666667, 0.0625,0.0588235, 0.0555556, 0.0526316}} may contain significant numerical errors."

Besides, I think trying to implement a simple solver is a good way for me as a new hand to learn the Mathematica language.

Answer 1

Here's a way to go about building the answer in a "loop free" manner. First define the matrix and vector:

 mat = {{1, 2, 0, 0, 0}, {2, 2, 3, 0, 0}, {0, 3, 3, 4, 0}, 
        {0, 0, 4, 4, 5}, {0, 0, 0, 5, 5}};;
 d = {5, 15, 31, 53, 45};
 {n, n} = Dimensions[mat];

Next extract out the a,b,c elements:

 b = Diagonal[mat];
 c = Flatten[{Diagonal[mat, 1], 0}];
 a = Flatten[{0, Diagonal[mat, -1]}];

Define three functions that we are going to iterate on, these are the same as the definitions of $c'$ and $d'$ and $x$ in your algorithm statement.

 cPrime[x_, {ax_, bx_, cx_}] := cx/(bx - x ax);
 dPrime[x_, {ax_, bx_, cpx_, dx_}] := (dx - x ax)/(bx - cpx ax);
 xOut[x_, {cpx_, dpx_}] := dpx - x cpx;

Now do the iterations using FoldList[ ]:

 cp = Rest[FoldList[cPrime, c[[1]]/b[[1]], Transpose[{a, b, c}]]];
 dp = Rest[FoldList[dPrime, d[[1]]/b[[1]], 
                  Transpose[{a, b, Flatten[{0, Drop[cp, -1]}], d}]]];
out = Reverse[FoldList[xOut, dp[[n]], 
                  Transpose[{Rest[Reverse[cp]], Rest[Reverse[dp]]}]]];

For the matrix you suggested, the output is {1,2,3,4,5} in agreement with your tridag[ ] module.

This can easily be made more concise using pure functions, which means that it is not necessary to define the cPrime, dPrime and xOut functions explicitly. In the pure notation form, the complete algorithm can be written:

 triThomas[a_, b_, c_, d_] := Module[{cp,dp}, 
    cp = Rest[FoldList[#2[[3]]/(#2[[2]] - #1 #2[[1]] ) &, c[[1]]/b[[1]], 
         Transpose[{a, b, c}]]];
    dp = Rest[FoldList[(#2[[4]] - #1 #2[[1]])/(#2[[2]] - #2[[3]] #2[[1]]) &, 
         d[[1]]/b[[1]], Transpose[{a, b, Flatten[{0, Drop[cp, -1]}], d}]]];
    Reverse[FoldList[(#2[[2]] - #1 #2[[1]]) &, dp[[n]], 
          Transpose[{Rest[Reverse[cp]], Rest[Reverse[dp]]}]]]]

Here, the symbol #1 refers to the first (implicit) argument and the symbols #2[[n]] refer to the nth element of the second argument.

This can be called by

 triThomas[a, b, c, d]

which gives the same answer as above.

Answer 2

Prior to version 5 users had to load an add on package that contained a tridiagonal solver based on the Thomas algorithm

I think the code for those old packages is probably accessible somewhere. I found this in the library archive:

http://library.wolfram.com/infocenter/MathSource/4827/

This package uses Compile. The earlier built in add-on package was actually quite slow from memory because it was procedural but not compiled. Users had the code though so could make their own compiled versions. Back in the day I had a functional version which was close to what @bill s has posted but reverted to the compiled procedural version for performance.