An example 2-d diffusion equation solver

Listed below is an example 2-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 2-d Poisson solver (see Sect. 5.10).

// CrankNicholson2D.cpp

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

//  dT / dt = D d^2 T / dx^2 +  D d^2 T / dy^2  for  xl <= x <= xh  
//                                              and  0 <= y <= L

//  alphaL T + betaL dT/dx = gammaL(y)  at x=xl

//  alphaH T + betaH dT/dx = gammaH(y)  at x=xh

// In y-direction, either simple Dirichlet boundary conditions:

//  T(x,0) = T(x,L) = 0

// or simple Neumann boundary conditions:

//  dT/dy(x,0) = dT/dy(x,L) = 0

// Matrix T assumed to be of extent N+2, J+1.
// Arrays gammaL, gammaH assumed to be of extent J+1.

// Now, (i,j)th elements of matrices correspond to

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

//  y_j = j * L / J      j=0,J

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

// Now, C =  D dt / dx^2, and kappa = pi * dx / L

// Finally, Neumann=0/1 selects Dirichlet/Neumann bcs in y-direction.

// Uses Crank-Nicholson scheme.

#include <blitz/array.h>

using namespace blitz;

void fft_forward_cos (Array<double,1> f, Array<double,1>& F);
void fft_backward_cos (Array<double,1> F, Array<double,1>& f);
void fft_forward_sin (Array<double,1> f, Array<double,1>& F);
void fft_backward_sin (Array<double,1> F, Array<double,1>& f);
void Tridiagonal (Array<double,1> a, Array<double,1> b, Array<double,1> c, 
                  Array<double,1> w, Array<double,1>& u);

void CrankNicholson2D (Array<double,2>& T,
                       double alphaL, double betaL, Array<double,1> gammaL, 
                       double alphaH, double betaH, Array<double,1> gammaH,
                       double dx, double C, double kappa, int Neumann)
{
  // Find N and J. Declare local arrays.
  int N = T.extent(0) - 2;
  int J = T.extent(1) - 1;
  Array<double,2> TT(N+2, J+1), V(N+2, J+1);
  Array<double,1> GammaL(J+1), GammaH(J+1);

  // Fourier transform T  
  for (int i = 0; i <= N+1; i++)
    {
      Array<double,1> In(J+1), Out(J+1);

      for (int j = 0; j <= J; j++) In(j) = T(i, j);

      if (Neumann)
        fft_forward_cos (In, Out);
      else
        fft_forward_sin (In, Out);

      for (int j = 0; j <= J; j++) TT(i, j) = Out(j);
    }

  // Fourier transform boundary conditions
  if (Neumann)
    {
      fft_forward_cos (gammaL, GammaL);
      fft_forward_cos (gammaH, GammaH);
    }
  else
    {
      fft_forward_sin (gammaL, GammaL);
      fft_forward_sin (gammaH, GammaH);
    }

  // Construct source term
  for (int i = 1; i <= N; i++)
    for (int j = 0; j <= J; j++) 
      V(i, j) = 
        0.5 * C * TT(i-1, j) + 
        (1. - C * (1. + 0.5 * double (j * j) * kappa * kappa)) * TT(i, j) +
        0.5 * C * TT(i+1, j);
     
  // Solve tridiagonal matrix equations
  if (Neumann)
    {
      for (int j = 0; j <= J; j++)
       {  
          Array<double,1> a(N+2), b(N+2), c(N+2), w(N+2), u(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 * (1. + 0.5 * double (j * j) * kappa * kappa);
          b(1) += 0.5 * C * betaL / (alphaL * dx - betaL);
          b(N) -= 0.5 * C * betaH / (alphaH * dx + betaH);
          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) = V(i, j);
          w(1) += 0.5 * C * GammaL(j) * dx / (alphaL * dx - betaL);
          w(N) += 0.5 * C * GammaH(j) * dx / (alphaH * dx + betaH);
	  
          // Invert tridiagonal matrix equation
          Tridiagonal (a, b, c, w, u);
          for (int i = 1; i <= N; i++) TT(i, j) = u(i);
       }
    }
  else
    {
      for (int j = 1; j < J; j++)
       { 
          Array<double,1> a(N+2), b(N+2), c(N+2), w(N+2), u(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 * (1. + 0.5 * double (j * j) * kappa * kappa);
          b(1) -= betaL / (alphaL * dx - betaL);
          b(N) += betaH / (alphaH * dx + betaH);
          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) = V(i, j);
          w(1) += 0.5 * C * GammaL(j) * dx / (alphaL * dx - betaL);
          w(N) += 0.5 * C * GammaH(j) * dx / (alphaH * dx + betaH);
	  
          // Invert tridiagonal matrix equation
          Tridiagonal (a, b, c, w, u);
          for (int i = 1; i <= N; i++) TT(i, j) = u(i);
       }
      
      for (int i = 1; i <= N  ; i++) 
        {
          TT(i, 0) = 0.; TT(i, J) = 0.;
        }
    }

  // Reconstruct solution
  for (int i = 1; i <= N; i++)
    {
      Array<double,1> In(J+1), Out(J+1);

      for (int j = 0; j <= J; j++) In(j) = TT(i, j);

      if (Neumann)
        fft_backward_cos (In, Out);
      else
        fft_backward_sin (In, Out);

      for (int j = 0; j <= J; j++) T(i, j) = Out(j);
    }

  // Calculate i=0 and i=N+1 values  
  for (int j = 0; j <= J; j++)
    {
      T(0, j) = (gammaL(j) * dx - betaL * T(1, j)) /
        (alphaL * dx - betaL);
      T(N+1, j) = (gammaH(j) * dx + betaH * T(N, j)) /
        (alphaH * dx + betaH);
    }
}