% A Tutorial of Geometric Multgrid Method for Solving Discretized Two-point ODE BVP.
% In the previous lecture, a Two-point ODE BVP is discretized by central finite difference method to get a sparse linear system, whose dimension grows as the mesh step size h=1/(m+1) goes to zero. To approximately solve such linear systems, iterative methods are more popular than direct method, since (1) iterative methods make best use of sparsity and (2) iterative methods can stop early whenever the approximation accuracy reaches the level of truncation errors due to discretization. The advantage of iterative methods becomes even significant for 2D and 3D cases.
% However, as the system dimension becomes larger, most standard iterative methods (e.g., Jacobi, Gauss-Seidel, SOR, CG, GMRES) show slower convergence rates, because the resulting coefficient matrix gets more ill-conditioned. Fortunately, Multgrid methods often can achieve the optimal mesh-independent convergence rates for certain class of systems. Once such a multgrid algorithm is developed, it can also be used as an effective precondtioner for other Krylov subspace solvers.
% Using the same example, we compare the iteration numbers of Jacobi and Gauss-Seidel iterative solvers as h\to 0.

c1=(3-exp(-2))/(exp(2)-exp(-2));c2=(exp(2)-3)/(exp(2)-exp(-2));
ufun=@(x) c1*exp(2*x)+c2*(exp(-2*x));%exact solution
ffun=@(x) zeros(size(x)); %right-hand size
mlist=2.^(3:7);hlist=1./(mlist+1);
tol=1e-6; maxit=100000; itVec=zeros(length(mlist),2);
for k=1:length(mlist)
    m=mlist(k); h=1/(m+1); x=linspace(0,1,m+2)'; e=ones(m,1);
    Ah=spdiags([e/h^2 -2*e/h^2-4 e/h^2],(-1:1),m,m);%tridiagonal matrix
    fh=ffun(x(2:end-1));%define rhs over interior nodes
    fh(1)=fh(1)-1/h^2;fh(end)=fh(end)-3/h^2;%enforce boundary condition
    %Jacobi iterations
    dinv = 1./diag(Ah);
    U0=rand(m,1);%random initial guess
    r0=norm(fh-Ah*U0);%initial residual norm
    for itJac=1:maxit
        U0=U0 + dinv.*(fh - Ah*U0);%Jacobi iteration
        rk=norm(fh-Ah*U0);%residual norm
        if(rk/r0<tol) %relative residual norm <tol
            break;
        end
    end
    %Gauss-Seidel iterations
    L = tril(Ah);
    U0=rand(m,1);%random initial guess
    r0=norm(fh-Ah*U0);%initial residual norm
    for itGS=1:maxit
        U0=U0 + L\(fh - Ah*U0);%Gauss-Seidel iteration
        rk=norm(fh-Ah*U0);%residual norm
        if(rk/r0<tol) %relative residual norm <tol
            break;
        end
    end
    itVec(k,:)=[itJac itGS]; %store both iteration numbers
end
semilogy(mlist,itVec,'-o');xlabel('dimension (m)'); ylabel('iteration numbers');
legend('Jacobi','Gauss-Seidel')

% For a fixed tolerance tol=1e-6, as the above figure shows, more iteration numbers are required as the dimension m is doubled. Remeber that the operation cost of each iteration is also doubled since the matrix size is also doubled. The iteration numbers are also too large for such small systems.
% Next, we will look into the short-term convergence behavior of the Jacobi iterations (also called relaxation or smoother). We fixed m=64 and choose the initial guess to include both low- and high-frequency components (as in the tutorial slide page 35).

m=64; h=1/(m+1); x=linspace(0,1,m+2)'; e=ones(m,1);
Ah=spdiags([e/h^2 -2*e/h^2-4 e/h^2],(-1:1),m,m);%tridiagonal matrix
fh=ffun(x(2:end-1));%define rhs over interior nodes
fh(1)=fh(1)-1/h^2;fh(end)=fh(end)-3/h^2;%enforce boundary condition
uh=ufun(x(2:end-1));%exact solution over interior nodes, for compute errors
%Jacobi iterations
dinv = 1./diag(Ah);
U0=sin(2*pi*(1:m)'/(m+1))+0.5*sin(16*pi*(1:m)'/(m+1))+0.5*sin(32*pi*(1:m)'/(m+1));%random initial guess
subplot(3,3,1);plot(x(2:end-1),U0-uh);title('Initial errors')
for j=1:8
    U0=U0 + dinv.*(fh - Ah*U0);%Jacobi iteration
    subplot(3,3,j+1);plot(x(2:end-1),U0-uh);title(sprintf('Errors after %d iter.',j))
end

% As expected, the initial errors become much smoother after the first few Jacobi iterations, where the high frequency parts (small waves) are reduced quickly while the low frequency parts (large waves) are almost the same. In some sense, such Jacobi iterations become less effective after the first few iterations since the errors become smooth.
% The Multigrid method performs only a few Jacobi iterations at different levels of meshes, where the smooth errors in the fine grid becomes non-smooth after being restricted to the coarse grid. At the coarse grid, the Jacobi iterations become effective again in reducing the high-frequency (non-smooth) errors. By doing this in a recursive manner, a few Jacobi iterations at different levels together can effective eliminate both low- and high-frequency errors, resulting a mesh-independent convergence rate. To build a complete multgrid algorihtm, we need to define (1) the smoother (Jacobi iteration) at each level, (2) coarse grid opeartor (matrix) at each level, (3) restriction operator, and (4) interpolation operator between fine and coarse grid.
% In the following codes, we test the convergence of a V-cycle multigrid algorithm based on Jacobi iterations as the smoother. Clearly, the used V-cycle iterations (=about 4 Jacobi iterations at fine grid)  are independent of mesh step size, and the errors should be reduced by 4 times as the mesh step size is halved. The CPU time should grows linearly with respect to m, which gives so-called optimal O(m) linear time complexity, since one needs at least O(m) operations to access it.

tol=1e-6;maxit=100;itVec2=[];
for fineL=3:7
    tic
    %Setup operators
    mg=[];
    for Level=fineL:-1:2
        [mg(Level).A]=genMatrix(2^Level); %generate system at each level
    end
    A=mg(fineL).A;
    %Setup right-hand-side
    xmin=0; xmax=1; N=2^fineL; h=(xmax-xmin)/N;
    xgrid=(0:N)'*h; %including boundary nodes
    b=ffun(xgrid); %rhs
    b(1)=ufun(xmin); b(end)=ufun(xmax);%boundary condition
    truex=ufun(xgrid); %true solution
    
    x=rand(size(b)); %initial guess
    r0=norm(b-A*x);rk=r0;
    err=[];res=[];
    %Begin Multigrid Algorithm: V-cycle
    for it=1:maxit
        x=mg_iter(mg,x,b,fineL,1,1,'JAC'); %V-cycle, smoother choices: JAC, GS 
        rk=norm(b-A*x);
        err(it)=max(abs(x-truex)); %max-norm error
        res(it)=rk/r0; %relative residual
        if(rk/r0<tol)
            break;
        end
    end
    tcpu=toc;
    fprintf('m=%d, V-cycle iter=%d, Error = %1.2e, Residual=%1.2e, CPU=%1.3f \n',N,it,err(end),res(end),tcpu)
    itVec2=[itVec2;it];
    subplot(1,2,1);semilogy([1:it],res,'-o');xlabel('iteration'); ylabel('residual norms');hold on
    subplot(1,2,2);semilogy([1:it],err,'-o');xlabel('iteration'); ylabel('error norms');hold on
end
figure
plot(mlist,itVec2,'-*'); xlabel('dimension (m)'); ylabel('iteration numbers');legend('Multigrid')

% Next, we show in more detail the codes for each component, that was used in the above V-cycle multgrid algorithm. Below are two quick tests of the restriction and interpolation operators.

rc=mg_restrict([1 5 3 9 7]) %test the restriction operator
rf=mg_interp([1 5 7])  %test the interpolation operator

%
% V-cycle multigrid algorithm: this is the main recursive V-cycle algorithm that runs pre- and post- Jacobi iterations.

function [x]=mg_iter(mg,x0,b,level,pre,post,smoother)
%Multigrid V-cycle Recursive code
A=mg(level).A; %mg is a cell of structures stored the system matrix of all levels
if(level==2) %Coarest level
    x=A\b;%direct solver, only a very small system
else
    x = mg_smooth(A,x0,b,pre,smoother); % presmooth
    r = b - A*x;%compute residual
    rc = mg_restrict(r);% Restrict residual
    %solve the Coarse grid system using V-cycle again, level goes down by 1
    cc = mg_iter(mg,zeros(size(rc)),rc,level-1,pre,post,smoother);%Recursive calling
    x = x + mg_interp(cc);% coarse grid correction
    x = mg_smooth(A,x,b,post,smoother);% postsmooth
end
end

% Coarse grid operator: for a given mesh dimension N, the system matrix is constructed at each level and stored.

function A=genMatrix(N)
xmin=0; xmax=1; h=(xmax-xmin)/N; e = ones(N-1,1);
%include boundary conditions as part of the system
A=[sparse(1,1,1,1,N+1);
    spdiags([e/h^2 -2*e/h^2-4 e/h^2], 0:2, N-1, N+1);%match with the previous system
    sparse(1,N+1,1,1,N+1)];
end

% Restriction operator: it restricts a vector rf on the fine grid to the coarse grid as rc, where full weight is usually better.

function [rc]=mg_restrict(rf)
N=length(rf)-1;
%rc=[rf(1);rf(3:2:N-1);rf(N+1)];%Injection
rc=[rf(1);(rf(2:2:N-2)+2*rf(3:2:N-1)+rf(4:2:N))/4;rf(N+1)];%Full weight, better
end

% Interpolation (Prologation) operator: it linearly interpolates a vector rc on the coarse grid to the fine grid as rf.

function [rf]=mg_interp(rc)
N=length(rc)-1;
rf=zeros(2*N+1,1);
rf(1:2:2*N+1)=rc;rf(2:2:2*N)=0.5*(rf(1:2:2*N-1)+rf(3:2:2*N+1));
end

%  smoother: the following function defines 3 different iterations, with initial guess x0 and only nv iterations. Notice that such a smoother itself can be used as iterative solver on the fixed level, which is often very slow for large system size.

function [x]=mg_smooth(A,x0,b,nv,smoother)
%smoother iterations (also called relaxations), only run nv iterations
x=x0;
switch smoother
    case 'JAC'
        w=2/3;%Weighted Jacobi Smoother, w=1 gives usual Jacobi
        dinv = 1./diag(A);
        for k = 1:nv
            x=x + w*dinv.*(b - A*x);
        end
    case 'GS'  %Gauss-Seidel Smoother
        L = tril(A);
        for k = 1:nv
            x=x + L\(b - A*x);
        end 
end
end

%