Tridiagonal Matrix Algorithm ("Thomas Algorithm") in C++

Tridiagonal Matrix Thomas Algorithm in C++

In the previous article on solving the heat equation via the Tridiagonal Matrix ("Thomas") Algorithm we saw how to take advantage of the banded structure of the finite difference generated matrix equation to create an efficient algorithm to numerically solve the heat equation. We will now provide a C++ implementation of this algorithm, and use it to carry out one timestep of solution in order to solve our diffusion equation problem.

The Tridiagonal Matrix Algorithm, also known as the Thomas Algorithm, is an application of gaussian elimination to a banded matrix. I've written up the mathematical algorithm in this article. The algorithm itself requires five parameters, each vectors. The first three parameters $\bf{a}$, $\bf{b}$ and $\bf{c}$ represent the elements in the tridiagonal bands. Since $\bf{b}$ represents the diagonal elements it is one element longer than $\bf{a}$ and $\bf{c}$, which represent the off-diagonal bands. The latter two parameters represent the solution vector $\bf{f}$ and $\bf{d}$, the right-hand column vector.

You can see the function prototype here:

void thomas_algorithm(const std::vector<double>& a,
                      const std::vector<double>& b,
                      const std::vector<double>& c,
                      const std::vector<double>& d,
                      std::vector<double>& f) {

Notice that $\bf{f}$ is a non-const vector, which means it is the only vector to be modified within the function.

It is possible to write the Thomas Algorithm function in a more efficient manner to override the $\bf{c}$ and $\bf{d}$ vectors. Finite difference methods require the vectors to be re-used for each time step, so the following implementation utilises two additional temporary vectors, $\bf{c^{*}}$ and $\bf{d^{*}}$. The memory for the vectors has been allocated within the function. Every time the function is called these vectors are allocated and deallocated, which is suboptimal from an efficiency point of view.

A proper "production" implementation would pass references to these vectors from an external function that only requires a single allocation and deallocation. However, in order to keep the function straightforward to understand I've not included this aspect:

std::vector<double> c_star(N, 0.0);
std::vector<double> d_star(N, 0.0);

The first step in the algorithm is to initialise the beginning elements of the $\bf{c^{*}}$ and $\bf{d^{*}}$ vectors:

c_star[0] = c[0] / b[0];
d_star[0] = d[0] / b[0];

The next step, known as the "forward sweep" is to fill the $\bf{c^{*}}$ and $\bf{d^{*}}$ vectors such that we form the second matrix equation $\bf{A^{*}}f=d^{*}$:

for (int i=1; i<N; i++) {
  double m = 1.0 / (b[i] - a[i] * c_star[i-1]);
  c_star[i] = c[i] * m;
  d_star[i] = (d[i] - a[i] * d_star[i-1]) * m;
}

Once the forward sweep is carried out the final step is to carry out the "reverse sweep". Notice that the vector $\bf{f}$ is actually being assigned here. The function itself is void, so we don't return any values.

for (int i=N-1; i-- > 0; ) {
  f[i] = d_star[i] - c_star[i] * d[i+1];
}

I've included a main function, which sets up the Thomas Algorithm to solve one time-step of the Crank-Nicolson finite difference method discretised diffusion equation. The details of the algorithm are not so important here, as I will be elucidating on the method in further articles when we come to solve the Black-Scholes equation. This is only provided in order to show you how the function works in a "real world" situation:

#include <cmath>
#include <iostream>
#include <vector>

// Vectors a, b, c and d are const. They will not be modified                                                                                                                                                        
// by the function. Vector f (the solution vector) is non-const                                                                                                                                                      
// and thus will be calculated and updated by the function.                                                                                                                                                          
void thomas_algorithm(const std::vector<double>& a,
                      const std::vector<double>& b,
                      const std::vector<double>& c,
                      const std::vector<double>& d,
                      std::vector<double>& f) {
  size_t N = d.size();

  // Create the temporary vectors                                                                                                                                                                                    
  // Note that this is inefficient as it is possible to call                                                                                                                                                         
  // this function many times. A better implementation would                                                                                                                                                         
  // pass these temporary matrices by non-const reference to                                                                                                                                                         
  // save excess allocation and deallocation                                                                                                                                                                         
  std::vector<double> c_star(N, 0.0);
  std::vector<double> d_star(N, 0.0);

  // This updates the coefficients in the first row                                                                                                                                                                  
  // Note that we should be checking for division by zero here                                                                                                                                                       
  c_star[0] = c[0] / b[0];
  d_star[0] = d[0] / b[0];

  // Create the c_star and d_star coefficients in the forward sweep                                                                                                                                                  
  for (int i=1; i<N; i++) {
    double m = 1.0 / (b[i] - a[i] * c_star[i-1]);
    c_star[i] = c[i] * m;
    d_star[i] = (d[i] - a[i] * d_star[i-1]) * m;
  }

  // This is the reverse sweep, used to update the solution vector f                                                                                                                                                 
  for (int i=N-1; i-- > 0; ) {
    f[i] = d_star[i] - c_star[i] * d[i+1];
  }
}

// Although thomas_algorithm provides everything necessary to solve                                                                                                                                                  
// a tridiagonal system, it is helpful to wrap it up in a "real world"                                                                                                                                               
// example. The main function below uses a tridiagonal system from                                                                                                                                                   
// a Boundary Value Problem (BVP). This is the discretisation of the                                                                                                                                                 
// 1D heat equation.                                                                                                                                                                                                 
int main(int argc, char **argv) {

  // Create a Finite Difference Method (FDM) mesh with 13 points                                                                                                                                                     
  // using the Crank-Nicolson method to solve the discretised                                                                                                                                                        
  // heat equation.                                                                                                                                                                                                  
  size_t N = 13;
  double delta_x = 1.0/static_cast<double>(N);
  double delta_t = 0.001;
  double r = delta_t/(delta_x*delta_x);

  // First we create the vectors to store the coefficients                                                                                                                                                           
  std::vector<double> a(N-1, -r/2.0);
  std::vector<double> b(N, 1.0+r);
  std::vector<double> c(N-1, -r/2.0);
  std::vector<double> d(N, 0.0);
  std::vector<double> f(N, 0.0);

  // Fill in the current time step initial value                                                                                                                                                                     
  // vector using three peaks with various amplitudes                                                                                                                                                                
  f[5] = 1; f[6] = 2; f[7] = 1;

  // We output the solution vector f, prior to a                                                                                                                                                                       
  // new time-step                                                                                                                                                                                                     
  std::cout << "f = (";
  for (int i=0; i<N; i++) {
    std::cout << f[i];
    if (i < N-1) {
      std:: cout << ", ";
    }
  }
  std::cout << ")" << std::endl << std::endl;

  // Fill in the current time step vector d                                                                                                                                                                          
  for (int i=1; i<N-1; i++) {
    d[i] = r*0.5*f[i+1] + (1.0-r)*f[i] + r*0.5*f[i-1];
  }

  // Now we solve the tridiagonal system                                                                                                                                                                             
  thomas_algorithm(a, b, c, d, f);

  // Finally we output the solution vector f                                                                                                                                                                         
  std::cout << "f = (";
  for (int i=0; i<N; i++) {
    std::cout << f[i];
    if (i < N-1) {
      std:: cout << ", ";
    }
  }
  std::cout << ")" << std::endl;

  return 0;
}

The output of the program is follows:

f = (0, 0, 0, 0, 0, 1, 2, 1, 0, 0, 0, 0, 0)

f = (0, 0, 0, 0.00614025, 0.145331, 0.99828, 1.7101, 0.985075, 0.143801, 0.0104494, 0.000759311, 0, 0)

You can see the the initial vector and then the solution vector after one time-step. Since we are modelling the diffusion/heat equation, you can see how the initial "heat" has spread out. Later on we will see how the diffusion equation represents the Black-Scholes equation with modified boundary conditions. One can think of solving the Black-Scholes equation as solving the heat equation "in reverse". In that process the "diffusion" happens against the flow of time. That will have to wait until a later article, however!