%
% Jacobi, GS and SOR for the 2-D Poisson problem
% discretized by five point finite difference scheme
%
clearvars;
close all;

tol = 1e-16;
maxit = 400;

n = 128;
h = 1/(n+1); % step size
hh = h*h;

u0 = zeros(n+2,n+2); % initial guess + boundary condition
x = 0:h:1;  y = 0:h:1;
u0(1,:) = 0.25 * (y .* y);
u0(n+2,:) = 0.25 * (1 + y .* y);
u0(:,1) = 0.25 * (x .* x);
u0(:,n+2) = 0.25 * (x .* x + 1);

f = -1; % f(x,y)
[X,Y] = meshgrid(x);
ut = 0.25*(X.*X + Y.*Y); % true solution
norm_ut = norm(ut(:));

iter = maxit; % number of iteration steps

% Jacobi iteration
relerr_Jacobi = zeros(iter+1,1); % relative error
relerr_Jacobi(1) = 1;

u = u0;
for k = 1 : iter
    v_old = u;
    for i = 2:n+1
        for j = 2:n+1
            u(i,j) = 0.25* (hh*f + v_old(i+1,j) + v_old(i-1,j) ...
                + v_old(i,j+1) + v_old(i,j-1));
        end
    end
    relerr_Jacobi(k+1) = norm(u(:) - ut(:)) / norm_ut;
end

% GS iteration
relerr_GS = zeros(iter+1,1); % relative error
relerr_GS(1) = 1;

u = u0;
for k = 1 : iter
    for i = 2:n+1
        for j = 2:n+1
            u(i,j) = 0.25* (hh*f + u(i+1,j) + u(i-1,j) ...
                                 + u(i,j+1) + u(i,j-1));
        end
    end
    relerr_GS(k+1) = norm(u(:) - ut(:)) / norm_ut;
end

% SOR iteration
relerr_SOR = zeros(iter+1,1); % relative error
relerr_SOR(1) = 1;

omega = 2 / (1+sin(pi/(n+1)));
u = u0;
for k = 1 : iter
    for i = 2:n+1
        for j = 2:n+1
            vij_old = u(i,j);
            u(i,j) = 0.25* (hh*f + u(i+1,j) + u(i-1,j) ...
                                 + u(i,j+1) + u(i,j-1));
            u(i,j) = (1-omega)*vij_old + omega*u(i,j);
        end
    end
    relerr_SOR(k+1) = norm(u(:) - ut(:)) / norm_ut;
end



A = gallery('poisson',n); % 2-D poisson equation
xt= ut(2:end-1,2:end-1); xt = xt(:);   % true solution
b = A*xt;           % right-hand side
x0 = zeros(n*n,1);

% CG iteration
[x, flag, relres, iter, relerr] = mycg(A,x0,b,tol,maxit,xt);

% plot the error
step = 10;
xx = [1,step+2 : step : iter+1];

semilogy(xx, relerr_Jacobi(xx),'-b+');
hold on
semilogy(xx, relerr_GS(xx),'-ro');
semilogy(xx, relerr_SOR(xx),'-kd');
hold on 
semilogy(xx,relerr(xx), 's:');
title(['n=',int2str(n)]);
set(gca,'FontSize',12)
legend('Jacobi', 'GS', 'SOR','CG');
% xlim([1,120]);
fprintf('||x-xt||/||xt|| = %.2e\n', norm(x-xt)/norm(xt));

%%%%% CG method %%%%%
function [x, flag, relres, iter, relerr] = mycg(A, x, b, tol, maxit, xt)
% 输入参数:
%   A: 系数矩阵
%   x: 迭代初值
%   b: 右端项
%   tol: 停机准则
%   maxit: 最大迭代步数
%
% 输出参数:
%   x: 近似解
%   flag: 算法成功标识
%   relres: 相对残量范数
%   iter: 迭代步数
%   resvec: 每次迭代的相对残量范数
%
flag = 0;
iter = 0;
relerr = zeros(maxit+1,1);
relerr(1) = 1;

% [dim,dim] = size(A);
bnrm2 = norm( b );
if  ( bnrm2 == 0.0 ), bnrm2 = 1.0; end

r = b - A*x;
relres = norm( r ) / bnrm2;
if ( relres < tol ) return, end

for iter = 1 : maxit
    rho = (r'*r);
    if ( iter > 1 )
        beta = rho / rho_1;
        p = r + beta*p;
    else
        p = r;
    end    
    q = A*p;
    alpha = rho / (p'*q );
    x = x + alpha * p;
    r = r - alpha*q;
    relres = norm( r ) / bnrm2;
    % resvec = [resvec; relres];
    relerr(iter+1) = norm(x-xt)/norm(xt);
    if ( relres <= tol ), break, end
    rho_1 = rho;
end
if ( relres > tol ) flag = 1; end
end


% % Jacobi iteration
% function [x, flag, relres, iter, resvec] = Jacobi(A, x, b, tol, maxit)
% 
% relerr_Jacobi = zeros(iter+1,1); % relative error
% relerr_Jacobi(1) = 1;
% 
% u = u0;
% for k = 1 : iter
%     v_old = u;
%     for i = 2:n+1
%         for j = 2:n+1
%             u(i,j) = 0.25* (hh*f + v_old(i+1,j) + v_old(i-1,j) ...
%                 + v_old(i,j+1) + v_old(i,j-1));
%         end
%     end
%     relerr_Jacobi(k+1) = norm(u(:) - ut(:)) / norm_ut;
% end

