% heat2d.m:  Solve the following IBVP for the 2D heat equation:
%
%            u_t = mu * (u_xx + u_yy)     for 0 <= x,y <= 1 
%            u(0,y,t) = u(1,y,t) = 0      (vertical sides held constant)
%            u_y(x,0,t) = u_y(x,1,t) = 0  (horizontal sides insulated)
%            u(x,y,0) = sin(3*pi*x) * cos(pi*y)
%
% This problem has exact solution:
%
%   u(x,y,t) = exp(-10*mu*pi^2*t) * sin(3*pi*x) * cos(pi*y)
%

tend= 1;    % 1 or 10
nx  = 20;
ny  = nx;
nt  = 50;
mu  = 0.1;

dx = 1 / nx;
dy = 1 / ny;
dt = tend / nt;
x = [ 0 : dx : 1 ]';
y = [ 0 : dy : 1 ]';
[xx,yy] = meshgrid(x,y); 

% Initial values.
umat = sin(3*pi*xx) .* cos(pi*yy);
% umat(1:ny+1,1)    = 1;  % alternate Dirichlet BC's
% umat(1:ny+1,nx+1) = 1;  % (uncomment these and 'uerr' is wrong)

% Set up the nn X nn coefficient matrix, 
% with nn=(nx-1)*(ny-1).
mx = mu*dt/dx^2;
my = mu*dt/dy^2;
e    = repmat( ones(nx-1,1), ny-1, 1);
esub = repmat( [ones(nx-2,1); 0], ny-1, 1);
esup = repmat( [0; ones(nx-2,1)], ny-1, 1);
nn = (nx-1)*(ny-1);
A = spdiags( [-my*e, -mx*esub, (1+2*mx+2*my)*e, -mx*esup, -my*e], ...
	     [-(nx-1), -1, 0, 1, nx-1], nn, nn );

% Fix up first (and last) nx-1 equations for Neumann BC's for u_y.
idx = sub2ind(size(A), [1:nx-1], [1:nx-1]);
A(idx) = 1+2*mx+my;                         
idx = sub2ind(size(A), [nn-nx+2:nn], [nn-nx+2:nn]);
A(idx) = 1+2*mx+my;  

spy(A), shg
fprintf(1, 'Enter to continue ...\n');
pause

% Time loop.  Some notes:
% 1. The solution and RHS are converted back and forth between matrix
%    and vector form.  This is necessary for plotting, and it is also 
%    more convenient to deal with boundary conditions when u is
%    stored as a matrix. 
% 2. Be careful of indexing in matrices.  Rows correspond to
%    different y-values, while columns to x-values, and so indexing
%    takes the form:  umat(y_index, x_index).

for i = 1 : nt,
  
  t = i*dt;
  
  % Extract the u values corresponding to the interior points. 
  rhsmat = umat(2:end-1, 2:end-1);
  
  % Update the Dirichlet boundary conditions (left/right sides). 
  rhsmat(1:ny-1,1)    = rhsmat(1:ny-1,1)    + mx*umat(2:ny,1);
  rhsmat(1:ny-1,nx-1) = rhsmat(1:ny-1,nx-1) + mx*umat(2:ny,nx+1);
  
  % Then reform the corresponding (ny-1) X (nx-1) matrix into a
  % vector.  Note: the transpose operation is required since
  % 'reshape' takes elements column-wise (we need row-wise).
  rhsvec = reshape( rhsmat', nn, 1 );
  
  % Solve the linear system.
  uvec = A \ rhsvec;

  % Reforms the solution vector into a matrix and assign to the
  % interior point values of the u matrix.  Note: the transpose
  % operation is required since 'reshape' takes elements column-wise.  
  umat(2:end-1,2:end-1) = reshape(uvec, nx-1, ny-1)';

  % Fix up the Neumann BC's (for plotting purposes only).
  umat(1,2:nx)    = umat(2,2:nx);   
  umat(ny+1,2:nx) = umat(ny,2:nx); 

  % Compute the exact solution and the point-wise error.
  uex = exp(-10*mu*pi^2*t) * sin(3*pi*xx) .* cos(pi*yy);
  uerr = umat - uex;
  
  % Plot the results and generate an animation of the solution.
  figure(1)
  surf( xx, yy, umat);
  set(gca,'ZLim',[-1,1])
  xlabel('x'), ylabel('y'), zlabel('u')
  title(['t=',num2str(t,'%6.2f')])
  grid on

  figure(2)
  contour( xx, yy, uerr );
  xlabel('x'), ylabel('y')
  title('Error')
  axis square
  colorbar
  grid on
  pause(0.1)
  
end
