function mg5ptdemo2(p)
%MG5PTDEMO2  demonstration of multigrid and multigrid preconditioning
%
% Demonstration of multigrid for 5-point Laplacian: shows error
% decrease for multigrid v-cycle and for conjugate gradients
% preconditioned by multigrid; meshsize is h = 1/2^p with a default
% value of 4 for p.
% Douglas N. Arnold, 2009-10-15

if nargin == 0
  p=4;
end
niters = 10;

% supply f
%f = @(x,y) -2*pi^2*sin(pi*x).*sin(pi*y);
f = @(x,y) -190*x.^2.*(1-y)+26.4*exp((x-.25).^2+(y-.25).^2);

% set mesh size
n = 2^p-1;
h = 1/(n+1);
x = linspace(0,1,n+2);   % x grid points including boundaries
y = linspace(0,1,n+2);   % y grid points including boundaries

[X,Y] = meshgrid(x,y);      % 2d arrays of x,y values
X = X';                     % transpose so that X(i,j),Y(i,j) are
Y = Y';                     % coordinates of (i,j) grid point, i.e., ((i-1)h,(j-1)h)

Iint = 2:n+1;                 % indices of interior points in x
Jint = 2:n+1;                 % indices of interior points in y
Xint = X(Iint,Jint);        % interior points
Yint = Y(Iint,Jint);

fh = f(Xint,Yint);          % fh is a mesh function on interior grid points

%%%%%%  Direct solve %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% set uh on the boundary to given values (interior values will be ignored)
uh = zeros(n+2,n+2);

% form 5-point Laplacian matrix (times h^2)
I = speye(n);
e = ones(n,1);
A = spdiags([-e,4*e,-e],[-1,0,1],n,n);
J = spdiags([-e,-e],[-1,1],n,n);
Lh = kron(I,A) + kron(J,I);

fh = f(Xint,Yint);          % fh is a mesh function on interior grid points
% multiply right hand side by h^2 and adjust for boundary condition
gh = h^2*fh;
gh(:,1) = gh(:,1) - uh(Iint,1);
gh(:,n) = gh(:,n) - uh(Iint,n+2);
gh(1,:) = gh(1,:) - uh(1,Jint);
gh(n,:) = gh(n,:) - uh(n+2,Jint);

% convert the right hand side from a grid function into a vector
F = reshape(gh,n*n,1);

% Solve the linear system for the grid function uh on the interior, written
% as vector
U = Lh\F;  

%%%%%%  End direct solve %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%% multigrid v-cycle %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
fprintf('Multigrid v-cycle for 5 point Laplacian\n\n')
% set up initial iterate (zero grid function)
u = zeros(size(fh));
% print max norm error
fprintf('iteration    0.   max norm error: %9.5f\n',max(abs(u(:)-U)))

% iterate and display error
for niter = 1:niters
  u = mg5pt(fh,u);
  % print max norm error
  fprintf('iteration %4d.   max norm error: %9.5f\n',niter,max(abs(u(:)-U)))
end
%%%%%%  End multigrid  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%% multigrid preconditioned CG %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
fprintf('\nConjugate gradients with multigrid preconditioner\n\n')
u = zeros(n);
% print max norm error
fprintf('iteration    0.   max norm error: %9.5f\n',max(abs(u(:)-U)))
% compute residual
r = fh;
s = mg5pt(r,u);
r2 = r(:)'*s(:);
for niter = 1:niter
  S = [zeros(1,n+2);zeros(n,1),s,zeros(n,1);zeros(1,n+2)];
  t = (4*s-S(1:n,2:n+1)-S(3:n+2,2:n+1)-S(2:n+1,1:n)-S(2:n+1,3:n+2))/h^2;
  s2 = s(:)'*t(:);
  lambda = r2/s2;
  u = u + lambda*s;
  % print max norm error
  fprintf('iteration %4d.   max norm error: %9.5f\n',niter,max(abs(u(:)-U)))
  r = r-lambda*t;
  r2old = r2;
  z = mg5pt(r,zeros(n));
  r2 = r(:)'*z(:);
  s=z+(r2/r2old)*s;
end
%%%%%%  End multigrid preconditioned CG %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%