# include <math.h>
# include <stdbool.h>
# include <stdio.h>
# include <stdlib.h>

# include "poisson_2d.h"

int main ( int argc, char *argv[] );
void jacobi ( int nx, int ny, double dx, double dy, double **f, double **u, 
  double **unew );

/******************************************************************************/

void poisson_2d ( int nx, int ny, void rhs ( int nx, int ny, double **f ), 
  double u_exact ( double x, double y ) )

/******************************************************************************/
/*
  Purpose:

    poisson_2d() solves the Poisson equation in the unit square.

  Discussion:

    This program runs serially.  Its output is used as a benchmark for
    comparison with similar programs run in a parallel environment.

    The Poisson equation

      - DEL^2 U(X,Y) = F(X,Y)

    is solved on the unit square [0,1] x [0,1] using a grid of NX by
    NX evenly spaced points.  The first and last points in each direction
    are boundary points.

    The boundary conditions and F are set so that the exact solution is

      U(x,y) = sin ( pi * x * y )

    so that

      - DEL^2 U(x,y) = pi^2 * ( x^2 + y^2 ) * sin ( pi * x * y )

    The Jacobi iteration is repeatedly applied until convergence is detected.

  Licensing:

    This code is distributed under the MIT license. 

  Modified:

    04 October 2024

  Author:

    John Burkardt
*/
{
  double a;
  double b;
  double c;
  bool converged;
  double d;
  double diff;
  double dx;
  double dy;
  double **f;
  double fnorm;
  int i;
  int it;
  int it_max = 30000;
  int j;
  double tolerance = 0.000001;
  double **u;
  double u_norm;
  double udiff;
  double uexact;
  double **unew;
  double unew_norm;
  double *x;
  double *y;
/*
  Print a message.
*/
  printf ( "\n" );
  printf ( "poisson_2d():\n" );
  printf ( "  A program for solving the Poisson equation.\n" );
  printf ( "\n" );
  printf ( "  -DEL^2 U = F(X,Y)\n" );
  printf ( "\n" );
  printf ( "  on the rectangle 0 <= X <= 1, 0 <= Y <= 1.\n" );
  printf ( "\n" );
  printf ( "  F(X,Y) = pi^2 * ( x^2 + y^2 ) * sin ( pi * x * y )\n" );
/*
  Set the initial solution estimate.
  We are "allowed" to pick up the boundary conditions exactly.
*/
  f = ( double ** ) malloc ( ny * sizeof ( double ) );
  for ( i = 0; i < ny; i++ )
  {
    f[i] = ( double * ) malloc ( nx * sizeof ( double ) );
  }

  u = ( double ** ) malloc ( ny * sizeof ( double ) );
  for ( i = 0; i < ny; i++ )
  {
    u[i] = ( double * ) malloc ( nx * sizeof ( double ) );
  }

  unew = ( double ** ) malloc ( ny * sizeof ( double ) );
  for ( i = 0; i < ny; i++ )
  {
    unew[i] = ( double * ) malloc ( nx * sizeof ( double ) );
  }

  a = 0.0;
  b = 1.0;
  dx = ( b - a ) / ( double ) ( nx - 1 );
  x = ( double * ) malloc ( nx * sizeof ( double ) );
  for ( j = 0; j < nx; j++ )
  {
    x[j] = ( ( nx - j ) * a + j * b ) / ( nx - 1 );
  }

  c = 0.0;
  d = 1.0;
  dy = ( d - c ) / ( double ) ( ny - 1 );
  y = ( double * ) malloc ( ny * sizeof ( double ) );
  for ( i = 0; i < ny; i++ )
  {
    y[i] = ( ( ny - i ) * c + i * d ) / ( ny - 1 );
  }

  printf ( "\n" );
  printf ( "  The number of interior X grid points is %d\n", nx );
  printf ( "  The number of interior Y grid points is %d\n", ny );
  printf ( "  The X grid spacing is %f\n", dx );
  printf ( "  The Y grid spacing is %f\n", dy );
/*
  Set the right hand side.
*/
  rhs ( nx, ny, f );
  fnorm = 0.0;
  for ( i = 0; i < ny; i++ )
  {
    for ( j = 0; j < nx; j++ )
    {
      fnorm = fnorm + pow ( f[i][j], 2 );
    }
  }
  fnorm = sqrt ( fnorm / nx / ny );
  printf ( "  RMS norm of F = %g\n", fnorm );
/*
  Initialize the solution estimate.
*/
  unew_norm = 0.0;
  for ( i = 0; i < ny; i++ )
  {
    for ( j = 0; j < nx; j++ )
    {
      if ( i == 0 || i == ny - 1 || j == 0 || j == nx - 1 )
      {
        unew[i][j] = f[i][j];
      }
      else
      {
        unew[i][j] = 0.0;
      }
      unew_norm = unew_norm + pow ( unew[i][j], 2 );
    }
  }
  unew_norm = sqrt ( unew_norm / nx / ny );
/*
  Set up the exact solution.
*/
  u_norm = 0.0;
  for ( i = 0; i < ny; i++ )
  {
    for ( j = 0; j < nx; j++ )
    {
      uexact = u_exact ( x[j], y[i] );
      u_norm = u_norm + pow ( uexact, 2 );
    }
  }
  u_norm = sqrt ( u_norm / nx / ny );
  printf ( "  RMS of exact solution = %g\n", u_norm );
/*
  Do the iteration.
*/
  converged = false;

  printf ( "\n" );
  printf ( "  Step    ||Unew||     ||Unew-U||     ||Unew-Exact||\n" );
  printf ( "\n" );

  udiff = 0.0;
  for ( i = 0; i < ny; i++ )
  {
    for ( j = 0; j < nx; j++ )
    {
      uexact = u_exact ( x[j], y[i] );
      udiff = udiff + pow ( unew[i][j] - uexact, 2 );
    }
  }
  udiff = sqrt ( udiff / nx / ny );
  printf ( "  %4d  %14g                  %14g\n", 0, unew_norm, udiff );

  for ( it = 1; it <= it_max; it++ )
  {
    for ( i = 0; i < ny; i++ )
    {
      for ( j = 0; j < nx; j++ )
      {
        u[i][j] = unew[i][j];
      }
    }
/*
  UNEW is derived from U by one Jacobi step.
*/
    jacobi ( nx, ny, dx, dy, f, u, unew );
/*
  Check for convergence.
*/
    u_norm = unew_norm;
 
    unew_norm = 0.0;
    diff = 0.0;
    udiff = 0.0;
    
    for ( i = 0; i < ny; i++ )
    {
      for ( j = 0; j < nx; j++ )
      {
        unew_norm = unew_norm + pow ( unew[i][j], 2 );
        diff = diff + pow ( unew[i][j] - u[i][j], 2 );
        uexact = u_exact ( x[j], y[i] );
        udiff = udiff + pow ( unew[i][j] - uexact, 2 );
      }
    }
    unew_norm = sqrt ( unew_norm / nx / ny );
    diff = sqrt ( diff / nx / ny );
    udiff = sqrt ( udiff / nx / ny );

    if ( it % 500 == 0 || diff <= tolerance )
    {
      printf ( "  %4d  %14g  %14g  %14g\n", it, unew_norm, diff, udiff );
    }
    if ( diff <= tolerance )
    {
      converged = true;
      break;
    }

  }

  if ( converged )
  {
    printf ( "  The iteration has converged.\n" );
  }
  else
  {
    printf ( "  The iteration has NOT converged.\n" );
  }
/*
  Free memory.
*/
  for ( i = 0; i < ny; i++ )
  {
    free ( f[i] );
    free ( u[i] );
    free ( unew[i] );
  }
  free ( f );
  free ( u );
  free ( unew );
  free ( x );
  free ( y );

  return;
}
/******************************************************************************/

void jacobi ( int nx, int ny, double dx, double dy, double **f, double **u, 
  double **unew )

/******************************************************************************/
/*
  Purpose:

   jacobi() carries out one step of the Jacobi iteration.

  Discussion:

    Assuming DX = DY, we can approximate

      - ( d/dx d/dx + d/dy d/dy ) U(X,Y) 

    by

      ( U(i-1,j) + U(i+1,j) + U(i,j-1) + U(i,j+1) - 4*U(i,j) ) / dx / dy

    The discretization employed below will not be correct in the general
    case where DX and DY are not equal.  It's only a little more complicated
    to allow DX and DY to be different, but we're not going to worry about 
    that right now.

  Licensing:

    This code is distributed under the MIT license. 

  Modified:

    26 October 2011

  Author:

    John Burkardt

  Input:

    int NX, NY, the X and Y grid dimensions.

    double DX, DY, the spacing between grid points.

    double **F, the right hand side data.

    double **U, the previous solution estimate.

  Output:

    double **UNEW, the updated solution estimate.
*/
{
  int i;
  int j;

  for ( i = 0; i < ny; i++ )
  {
    for ( j = 0; j < nx; j++ )
    {
      if ( i == 0 || j == 0 || i == ny - 1 || j == nx - 1 )
      {
        unew[i][j] = f[i][j];
      }
      else
      { 
        unew[i][j] = 0.25 * ( 
          u[i-1][j] + u[i][j+1] + u[i][j-1] + u[i+1][j] + f[i][j] * dx * dy );
      }
    }
  }
  return;
}