%Solve 2D Poisson equation with Geometric Multgrid method
clear
global LEVEL0;
fprintf('N\t Rel. Res.\t Error \t iter \t CPU \n');
maxit=50;tol=1e-8;
for n=5:10
tic
fineL=n;
xmin=0;xmax=1;
ymin=0;ymax=1;
%Setup operators
mg=[];

LEVEL0=2;
%define the matrix at each level
for Level=fineL:-1:LEVEL0
n=2^Level; h=(xmax-xmin)/n;
[mg(Level).A]=(1/h^2)*gallery('poisson',n-1);%genMatrix(n,h);
end

A=mg(fineL).A;
%Setup right-hand-side
N=2^fineL;
hx=(xmax-xmin)/N;
hy=(ymax-ymin)/N;
xgrid=xmin+(1:N-1)'*hx;%interior nodes
ygrid=ymin+(1:N-1)'*hy;
[X,Y] = meshgrid(xgrid,ygrid);
S0=U0(X,Y);%true solution
%For original equaion
F=fF(X,Y);
%F(1,:)=0;F(end,:)=0;F(:,1)=0;F(:,end)=0;
b=F(:);
%Begin Multigrid Algorithm
x=zeros(size(b));
h=hx;
r0=h*norm(b-A*x);rk=r0;
it=1;
%x=full_mg_2d(mg,x,b,fineL,1,1,'GS');%Full Multigrid, faster
while((it<=maxit)&&(rk>=r0*tol))
x=mg_iter_2d(mg,x,b,fineL,2,1,'GS');
rk=h*norm(b-A*x);
err0(it)=rk/r0;
it=it+1;
end
iter=it-1;%total V-cycle iterations
semilogy([1:it-1],(err0),'-o');hold on
res=err0(end);
err=norm(S0(:)-x,inf);%solution max error, decrease by 4 times
cpu=toc;%CPU time should grows linearly, 4 times
fprintf('%d&\t %1.1e & \t %1.1e &\t %d&\t %1.2f \\\\ \n',N,res,err,iter,cpu);

end

% function [A]=genMatrix(N,h)
% e = (1/h^2)*ones(N-1,1);
% E=speye(N-1);
% T=spdiags([e -2*e e], -1:1, N-1, N-1);
% A=kron(T,E)+kron(E,T);
% end

function [x]=mg_iter_2d(mg,x0,b,level,pre,post,smoother)
%fprintf('+++++++++++++++++++Level[%d] Begin+++++++++++++++++++++\n',level);
A=mg(level).A;
global LEVEL0;
if(level==LEVEL0) %Coarest level
x=A\b;
else
% presmooth
%fprintf('Level[%d],Before Presmooth res=%1.2e\n',level,norm(b-A*x0));
x = mg_smooth(A,x0,b,pre,smoother);
%fprintf('Level[%d],After Presmooth res=%1.2e\n',level,norm(b-A*x));
% Restrict residual
r = b - A*x;
%rc = mg_restricth(r,h);
rc = rest(r);
% coarse grid correction
cc = mg_iter_2d(mg,zeros(size(rc)),rc,level-1,pre,post,smoother);
%cc = mg_iter(mg,cc,rc,level-1,pre,post,smoother);
x = x + intp2(cc);
%x=x + mg_interph(cc,h/2);
% postsmooth
%fprintf('Level[%d],Before Postsmooth res=%1.2e\n',level,norm(b-A*x));
x = mg_smooth(A,x,b,post,smoother);
%fprintf('Level[%d],After Postsmooth res=%1.2e\n',level,norm(b-A*x));
end
%fprintf('+++++++++++++++++++Level[%d] End+++++++++++++++++++++\n',level);
end
function [x]=mg_smooth(A,x0,b,nv,smoother)
x=x0;
%Jacobi Smoother
switch smoother
case 'JAC'
w=2/3;
dinv = 1./diag(A);
for k = 1:nv
x=x + w*dinv.*(b - A*x);
end

case 'GS' %G-S Smoother
L = tril(A);
for k = 1:nv
x=x + (L\(b - A*x));
end
case 'RB-GS'
dinv = 1./diag(A);
for k = 1:nv
x(1:2:end)=x(1:2:end) + dinv(1:2:end).*(b(1:2:end) - A(1:2:end,:)*x);
x(2:2:end)=x(2:2:end) + dinv(2:2:end).*(b(2:2:end) - A(2:2:end,:)*x);
end
end
end
%
function [rc]=rest(r)
%restriction
n=sqrt(length(r));
r=reshape(r,n,n);
%half-weighting restriction
% rc=( +r(1:2:end-2,2:2:end-1) +...
% r(2:2:end-1,1:2:end-2)+4*r(2:2:end-1,2:2:end-1)+r(2:2:end-1,3:2:end)+...
% +r(3:2:end,2:2:end-1) )/8;
% rc=rc(:);
%full weighting averaging
rc=(r(1:2:end-2,1:2:end-2)+2*r(1:2:end-2,2:2:end-1) +r(1:2:end-2,3:2:end)+...
2*r(2:2:end-1,1:2:end-2)+4*r(2:2:end-1,2:2:end-1)+2*r(2:2:end-1,3:2:end)+...
r(3:2:end,1:2:end-2) +2*r(3:2:end,2:2:end-1) +r(3:2:end,3:2:end))/16;
rc=rc(:);
end

function rf=intp2(r)
%interpolation
n=sqrt(length(r));
R=zeros(n+2,n+2);
R(2:end-1,2:end-1)=reshape(r,n,n);
m=2*n+1;
rf=zeros(m,m);
rf(2:2:end-1,2:2:end-1) = R(2:end-1,2:end-1);
rf(1:2:end,2:2:end-1) = 0.5*(R(1:end-1,2:end-1) + R(2:end,2:end-1));
rf(2:2:end,1:2:end) = 0.5*(R(2:end-1,1:end-1) + R(2:end-1,2:end));
rf(1:2:end,1:2:end) = 0.25*(R(1:end-1,1:end-1)+ R(2:end,1:end-1)...
+ R(1:end-1,2:end) + R(2:end,2:end));
rf=rf(:);

end

function z=fF(x,y)
z=2*pi^2*sin(pi*x).*sin(pi*y);
%z=2*(1-6*x.^2).*(y.^2-y.^4)+2*(1-6*y.^2).*(x.^2-x.^4);
end
function z=U0(x,y)
z=sin(pi*x).*sin(pi*y);
%z=(x.^2-x.^4).*(y.^4-y.^2);
end

function [x]=full_mg_2d(mg,x,b,level,pre,post,smoother)
global LEVEL0;
A=mg(level).A;
r=b-A*x;
rc=r;
for i=level:-1:LEVEL0+1
rc=rest(rc);
end
zc=mg(LEVEL0).A\rc;
for i=LEVEL0:level-1
zf=intp2(zc);
rc=r;
for j=level:-1:i+2
rc=rest(rc);
end
zc=mg_iter_2d(mg,zf,rc,i+1,pre,post,smoother);
end
x=x+zc;
end