function exercise5 ( )

%*****************************************************************************80
%
%% exercise5: Crank-Nicolson code for 2D heat equation solvers.
%
%  Licensing:
%
%    This code is distributed under the GNU LGPL license.
%
%  Modified:
%
%    16 July 2020
%
%  Author:
%
%    John Burkardt
%
  fprintf ( 1, '\n' );
  fprintf ( 1, 'exercise5:\n' );
  fprintf ( 1, '  Use the Crank-Nicolson method for a 2D heat equation.\n' );
%
%  Discretization.
%
  nx = 11;
  a = -1.0;
  b = +1.0;
  x = linspace ( a, b, nx );
  dx = ( b - a ) / ( nx - 1 );

  ny = 11;
  c = -1.0;
  d = +1.0;
  y = linspace ( c, d, ny );
  dy = ( d - c ) / ( ny - 1 );

  nt = 11;
  e = 0.0;
  f = 1.0;
  t = linspace ( e, f, nt );
  dt = ( f - e ) / ( nt - 1 );

  kappa = 1.0 / pi^2;
%
%  Create some arrays.
%
  [ X, Y ] = ndgrid ( x, y );
  U = zeros(nx,ny);
  A = zeros(nx*ny,nx*ny);  
  rhs = zeros(nx*ny,1);
%
%  Do the timesteps.
%
  for k = 1 : nt
%
%  First step just sets the initial condition.
%
    if ( k == 1 )

      U = g ( X, Y, t(k) );

      umax = max ( max ( U ) );
      umin = min ( min ( U ) );
%
%  Later steps do the forward Euler update.
%
    else

      for j = 1 : ny
        for i = 1 : nx
%
%  Get the index C for node(i,j), and its NSEW neighbors.
%
          [ C, N, S, E, W ] = cnsew ( i, j, nx );
%
%  If on the boundary, set U = boundary value.
%
          if ( j == 1 || j == ny || i == 1 || i == nx )

            A(C,C) = 1.0;
            rhs(C) = g ( x(i), y(j), t(k) );
%
%  The Crank-Nicolson linear system can be written as the "average"
%  of the forward and backward Euler methods.
%
          else

            fw = 1.0;
            bw = 1.0 + 2.0 * kappa * dt / dx^2 + 2.0 * kappa * dt / dy^2;
            A(C,C) = 0.5 * fw + 0.5 * bw;

            fw = 0.0;
            bw = - kappa * dt / dx^2;
            A(C,W) = 0.5 * fw + 0.5 * bw;

            fw = 0.0;
            bw = - kappa * dt / dx^2;
            A(C,E) = 0.5 * fw + 0.5 * bw;

            fw = 0.0;
            bw = - kappa * dt / dy^2;
            A(C,N) = 0.5 * fw + 0.5 * bw;
                    
            fw = 0.0;
            bw = - kappa * dt / dy^2;        
            A(C,S) = 0.5 * fw + 0.5 * bw;
                                          
            fw = U(C) + kappa * dt/dx^2 * ( + U(E) - 2.0 * U(C) + U(W) ) ...
                      + kappa * dt/dy^2 * ( + U(N) - 2.0 * U(C) + U(S) );
            bw = U(C);
            rhs(C) = 0.5 * fw + 0.5 * bw;

          end

        end
      end
%
%  Solve the discretized linear system and reshape u as an array U.
%  
      u = A \ rhs;
      U = reshape ( u, nx, ny );

    end
%
%  Since we know the exact solution, we can print the RMS error.
%
    V = g ( X, Y, t(k) );
    E = rms ( reshape ( U - V, nx*ny, 1 ) );
    fprintf ( 1, '  Time %g, RMS error = %g\n', t(k), E );
%
%  Do a surface plot of U.
%
    surf ( X, Y, U );
    label = sprintf ( 'Crank-Nicolson Time step %d', k );
    axis ( [ -1.0, +1.0, -1.0, +1.0, umin, umax ] );
    title ( label );
    colorbar ( );
    pause ( );

  end

  return
end

function [ C, N, S, E, W ] = cnsew ( i, j, nx )

%*****************************************************************************80
%
%% cnsew returns the variable indices of the 5 point Laplacian stencil.
%
%  Input:
%
%    integer i, j: the grid indices of a point.
%
%    integer nx: the number of nodes in the x direction.
%
%  Output:
%
%    integer C, N, S, E, W: the variable indices for the central,
%    north, south, east, and west nodes.
%
  C = i     + ( j - 1 ) * nx;
  N = i     + ( j     ) * nx;
  S = i     + ( j - 2 ) * nx;
  E = i + 1 + ( j - 1 ) * nx;
  W = i - 1 + ( j - 1 ) * nx;

  return
end

function value = g ( x, y, t )

%*****************************************************************************80
%
%% g returns the value of the exact solution.
%
%  Discussion:
%
%    This function satisfies a form of the heat equation:
%
%      gt = - 1/pi^2 * gxx
%
  value = sin ( 3.0 * pi * x / 5.0 ) .* cos ( 4.0 * pi * y / 5.0 ) * exp ( - t );

  return
end

