Next: An example 2-d solution Up: The diffusion equation Previous: 2-d problem with Neumann

## 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);
}
}
```

Richard Fitzpatrick 2006-03-29