function [ errsor, errcg, errcgp ] = cgiter ( N, nit, iplot );
% CGIter
%
% This m-file is part of the material of the lectures on
%   Numerical Mathematics
% at the Math Department of Halle University, Germany.
%
% Author :      Prof. Dr. M. Arnold, arnold@mathematik.uni-halle.de
% Version of :  Jul 13, 2005
%
% Visualization of CG methods for the solution of linear systems
%
% Parameters:
%   N      (input)  : number of grid points in FD approximation
%   nit    (input)  : number of iteration steps
%   iplot  (input)  : control flag "graphical output"
%                       =0 .. no graphical output
%                       =1 .. 3D animation
%   errsor (output) : error of SOR iteration with optimal \omega (discrete l_2 norm)
%   errcg  (output) : error of classical CG method (discrete l_2 norm)
%   errcgp (output) : error of preconditioned CG method (discrete l_2 norm)
%
% Examples:
%
%   [ errsor, errcg, errcgp ] = cgiter ( 160, 100, 1 );
%
%   see also plotiter.m
%
%   Reference: Section 8.7 of
%   Stoer, J.; Bulirsch, R.: Numerische Mathematik 2, Dritte Auflage.
%   Springer-Verlag Berlin Heidelberg New York, 1990.

% -> finite difference approximation of Poisson problem in 2D
h = 1/N;
x = linspace ( 0, 1, N+1 );
[ XX, YY ] = meshgrid ( x );

% -> exact solution
e  = ones ( (N-1)^2, 1 );
ee = ones ( (N-1)^2, 1 );
ee(N-1:N-1:(N-1)^2) = 0;
L = spdiags ( [ -e -ee ] / h^2, [ -(N-1) -1 ], (N-1)^2, (N-1)^2 );
A = L + spdiags ( 4*e / h^2, 0, (N-1)^2, (N-1)^2 ) + L';
b = - ones( (N-1)^2, 1 );

xsol = A \ b;
Usol = zeros ( N+1, N+1 );
for i=1:N-1,
    Usol(1+(1:N-1),1+i) = xsol((i-1)*(N-1)+(1:N-1));
end;

umax = max ( max ( -Usol ) );

% -> initialize iteration
errsor = zeros ( nit+1, 1 );
errcg  = zeros ( nit+1, 1 );
errcgp = zeros ( nit+1, 1 );

Usor = zeros ( N+1, N+1 );
Ucg  = zeros ( N+1, N+1 );
Ucgp = zeros ( N+1, N+1 );

omega = 2 / ( 1 + sin ( pi / N ) );

rcg = b;
pcg = rcg;
ncgold = rcg' * rcg;

rcgp = - b;
qcgp = - pcgssor ( N, omega, rcgp );
pcgp = qcgp;
ncgpold = rcgp' * qcgp;

% -> iteration
for it=0:nit,
    
    % -> evaluate error of actual iterates
    errsor(1+it) = sqrt ( sum ( sum ( ( Usor-Usol ).^2 ) ) / (N-1)^2 ); 
    errcg(1+it)  = sqrt ( sum ( sum ( ( Ucg -Usol ).^2 ) ) / (N-1)^2 ); 
    errcgp(1+it) = sqrt ( sum ( sum ( ( Ucgp-Usol ).^2 ) ) / (N-1)^2 ); 
    
    % -> visualization (if applicable)
    if iplot==1,
        for isub=1:4,
            subplot ( 2, 2, isub );  
            if isub==1,
                surf ( XX, YY, -Ucg  );  title ( 'CG method' );
            elseif isub==2,
                surf ( XX, YY, -Ucgp );  title ( 'CG (preconditioned)' );
            elseif isub==3,
                surf ( XX, YY, -Usor );  title ( 'SOR method' );
            elseif isub==4,
                surf ( XX, YY, -Usol );  title ( 'Exact solution' );
            end;    
            xlabel ( 'x' );  ylabel ( 'y' );  zlabel ( '- u(x,y)' );
            set ( gca, 'ZLim', [ 0, 1.0*umax ] );
            set ( gca, 'CLim', [ 0,     umax ] );
            shading interp;
        end;
        set ( gcf, 'Name', sprintf ( 'Iteration step   k = %i', it ) );
        drawnow,
    end;
    
    % -> iteration (max. nit iteration steps)
    if it<nit,
        
        % -> SOR
        for i=1:N-1,  for j=1:N-1,
            Usor(1+i,1+j) = ( 1 - omega ) * Usor(1+i,1+j)             ...
                           + omega * ( -h^2                           ...
                           + Usor(1+(i-1),1+j) + Usor(1+(i+1),1+j)    ...
                           + Usor(1+i,1+(j-1)) + Usor(1+i,1+(j+1)) ) / 4;
        end;  end;

        % -> CG (classical)
        xcg = zeros ( (N-1)^2, 1 );
        for i=1:N-1,
            xcg((i-1)*(N-1)+(1:N-1)) = Ucg(1+(1:N-1),1+i);
        end;
        
        if norm(pcg)>10*eps,        
            Apcg = A * pcg;
            acg  = ncgold / ( pcg' * Apcg );
            xcg  = xcg + acg * pcg;
            rcg  = rcg - acg * Apcg;
        
            ncgnew = rcg' * rcg;
            bcg    = ncgnew / ncgold;
            ncgold = ncgnew;
        
            pcg = rcg + bcg * pcg;            
        end;
        
        Ucg = zeros ( N+1, N+1 );
        for i=1:N-1,
            Ucg(1+(1:N-1),1+i) = xcg((i-1)*(N-1)+(1:N-1));
        end;

        % -> CG (preconditioned)
        xcgp = zeros ( (N-1)^2, 1 );
        for i=1:N-1,
            xcgp((i-1)*(N-1)+(1:N-1)) = Ucgp(1+(1:N-1),1+i);
        end;
        
        if norm(pcgp)>10*eps,    
            Apcgp = - A * pcgp;
            acgp  = ncgpold / ( pcgp' * Apcgp );
            xcgp  = xcgp + acgp * pcgp;
            rcgp  = rcgp - acgp * Apcgp;
            
            qcgp = - pcgssor ( N, omega, rcgp );
        
            ncgpnew = rcgp' * qcgp;
            bcgp    = ncgpnew / ncgpold;
            ncgpold = ncgpnew;
        
            pcgp = qcgp + bcgp * pcgp;            
        end;
        
        Ucgp = zeros ( N+1, N+1 );
        for i=1:N-1,
            Ucgp(1+(1:N-1),1+i) = xcgp((i-1)*(N-1)+(1:N-1));
        end;
        
    end;
    
end;    
