Next: An example solution of Up: Poisson's equation Previous: The fast Fourier transform

## An example 2-d Poisson solving routine

Listed below is an example 2-d Poisson solving routine which employs the previously listed tridiagonal matrix inversion and FFT wrapper routines, as well as the Blitz++ library.
```// Poisson2d.cpp

// Function to solve Poisson's equation in 2-d:

//  d^2 u / dx^2 + d^2 u / dy^2 = v  for  xl <= x <= xh  and  0 <= y <= L

//  alphaL u + betaL du/dx = gammaL(y)  at x=xl

//  alphaH u + betaH du/dx = gammaH(y)  at x=xh

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

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

// or simple Neumann boundary conditions:

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

// Matrices u and v 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, kappa = pi * dx / L

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

#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 Poisson2D (Array<double,2>& u, Array<double,2> v,
double alphaL, double betaL, Array<double,1> gammaL,
double alphaH, double betaH, Array<double,1> gammaH,
double dx, double kappa, int Neumann)
{
// Find N and J. Declare local arrays.
int N = u.extent(0) - 2;
int J = u.extent(1) - 1;
Array<double,2> V(N+2, J+1), U(N+2, J+1);
Array<double,1> GammaL(J+1), GammaH(J+1);

// 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);
}

// Fourier transform source term
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) = v(i, j);

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

for (int j = 0; j <= J; j++) V(i, j) = Out(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), uu(N+2);

// Initialize tridiagonal matrix
for (int i = 2; i <= N; i++) a(i) = 1.;
for (int i = 1; i <= N; i++)
b(i) = -2. - 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) = 1.;

// Initialize right-hand side vector
for (int i = 1; i <= N; i++)
w(i) = V(i, j) * dx * dx;
w(1) -= GammaL(j) * dx / (alphaL * dx - betaL);
w(N) -= GammaH(j) * dx / (alphaH * dx + betaH);

// Invert tridiagonal matrix equation
Tridiagonal (a, b, c, w, uu);
for (int i = 1; i <= N; i++) U(i, j) = uu(i);
}
}
else
{
for (int j = 1; j < J; j++)
{
Array<double,1> a(N+2), b(N+2), c(N+2), w(N+2), uu(N+2);

// Initialize tridiagonal matrix
for (int i = 2; i <= N; i++) a(i) = 1.;
for (int i = 1; i <= N; i++)
b(i) = -2. - 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) = 1.;

// Initialize right-hand side vector
for (int i = 1; i <= N; i++)
w(i) = V(i, j) * dx * dx;
w(1) -= GammaL(j) * dx / (alphaL * dx - betaL);
w(N) -= GammaH(j) * dx / (alphaH * dx + betaH);

// Invert tridiagonal matrix equation
Tridiagonal (a, b, c, w, uu);
for (int i = 1; i <= N; i++) U(i, j) = uu(i);
}

for (int i = 1; i <= N  ; i++)
{
U(i, 0) = 0.; U(i, J) = 0.;
}
}

// Reconstruct solution via inverse Fourier transform
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) = U(i, j);

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

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

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

Richard Fitzpatrick 2006-03-29