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