m=40;

%set up Laplace equation
h = 1.0/(m+1);
I = speye(m,m);
E = sparse(2:m,1:m-1,1,m,m);
D = 2*I-E-E';
A = (kron(D,I)+kron(I,D))/h^2;
b = ones(m*m,1); %some arbitrary right-hand-side

xtrue = A\b; %true solution


diagonal = diag(diag(A));
L = tril(A);
U = A-L;
%for SOR
sorL = tril(A,-1);
%TODO clean up and make the notation consistent :-)

omega = 2 / (1+sin(pi*h)); %optimal choice

xgs = 100*rand(m*m,1); %initial guess with random values
xsor = xgs;
xcg = xgs;
xpcg = xgs;

stepcount = 20;
errnormsor = zeros(stepcount,1);
errnormcg = zeros(stepcount,1);
errnormpcg = zeros(stepcount,1);

%form incomplete Cholesky factor L
L=ichol(A);

for i = 1:stepcount
    
    %store the error for each step
    errnormsor(i) = norm(xsor-xtrue);
    errnormcg(i) = norm(xcg-xtrue);
    errnormpcg(i) = norm(xpcg-xtrue);
    
    %solvers
    xsor = (diagonal + omega * sorL) \ (omega*b - (omega*U + (omega-1)*diagonal) * xsor);
    
    %this is not ideal, since this way the CG solver starts from scratch each
    %iteration! This is only to be able to plot solution for every iteration
    xcg = pcg(A,b, 1e-17,i,eye(m*m),eye(m*m),xgs); %unpreconditioned 
    xpcg = pcg(A,b,1e-17,i,L,L',xgs); %precondition with IC
    
    %show plots
    figure(1);
    subplot(2,2,1);
    height = reshape(xcg-xtrue,m,m);
    bar3(height);
    zlim([0, 100]);
    title('CG');
    
    subplot(2,2,2);
    height = reshape(xpcg-xtrue,m,m);
    bar3(height);
    zlim([0, 100]);
    title('PCG');
    
    subplot(2,2,3);
    height = reshape(xsor-xtrue,m,m);
    bar3(height);
    zlim([0, 100]);
    title('optimal SOR');
    
    pause(0.1);
end

figure(2);
clf;
plot(errnormsor);
hold on;
plot(errnormcg);
plot(errnormpcg);
legend('SOR','CG','PCG');
ylim([0,2600]);
title('Soln error vs. iteration count');
