An improved 1-d diffusion equation solver

Listed below is an improved 1-d diffusion equation solver which uses the Crank-Nicholson scheme, as well as the previous listed tridiagonal matrix solver and the Blitz++ library. Note the great structural similarity between this solver and the previously listed 1-d Poisson solver (see Sect. 5.5).

// CrankNicholson1D.cpp

// Function to evolve diffusion equation in 1-d:

//  dT / dt = D d^2 T / dx^2  for  xl <= x <= xh

//  alpha_l T + beta_l dT/dx = gamma_l  at x=xl

//  alpha_h T + beta_h dT/dx = gamma_h  at x=xh

// Array T assumed to be of extent N+2.

// Now, ith element of array corresponds to

//  x_i = xl + i * dx    i=0,N+1

// Here, dx = (xh - xl) / (N+1) is grid spacing.

// Function evolves T by single time-step.

// C = D dt / dx^2, where dt is time-step.

// Uses Crank-Nicholson implicit scheme.

#include <blitz/array.h>

using namespace blitz;

void Tridiagonal (Array<double,1> a, Array<double,1> b, Array<double,1> c, 
                  Array<double,1> w, Array<double,1>& u);

void CrankNicholson1D (Array<double,1>& T, 
                  double alpha_l, double beta_l, double gamma_l,
                  double alpha_h, double beta_h, double gamma_h,
                  double dx, double C)
{
  // Find N. Declare local arrays.
  int N = T.extent(0) - 2;
  Array<double,1> a(N+2), b(N+2), c(N+2), w(N+2);

  // Initialize tridiagonal matrix
  for (int i = 2; i <= N; i++) a(i) = - 0.5 * C;
  for (int i = 1; i <= N; i++) b(i) = 1. + C;
  b(1) += 0.5 * C * beta_l / (alpha_l * dx - beta_l);
  b(N) -= 0.5 * C * beta_h / (alpha_h * dx + beta_h);
  for (int i = 1; i <= N-1; i++) c(i) = - 0.5 * C;

  // Initialize right-hand side vector
  for (int i = 1; i <= N; i++)
    w(i) = T(i) + 0.5 * C * (T(i-1) - 2. * T(i) + T(i+1));
  w(1) += 0.5 * C * gamma_l * dx / (alpha_l * dx - beta_l);
  w(N) += 0.5 * C * gamma_h * dx / (alpha_h * dx + beta_h);

  // Invert tridiagonal matrix equation
  Tridiagonal (a, b, c, w, T);

  // Calculate i=0 and i=N+1 values
  T(0) = (gamma_l * dx - beta_l * T(1)) /
    (alpha_l * dx - beta_l);
  T(N+1) = (gamma_h * dx + beta_h * T(N)) /
    (alpha_h * dx + beta_h);
}