
clear all
grid_no  = input('mesh number = ');
% x        = linspace(0, 1, grid_no);
% y        = linspace(0, 1, grid_no);
leng_h   = 1 / ( grid_no + 1 );
x        = leng_h:leng_h:grid_no*leng_h;
y        = x;
[xx, yy] = meshgrid(x, y);
zz       = sin(pi*xx).*sin(pi*yy);

tic;
%
% Generate the coefficient matrix of the linear system
%
i_idx = zeros(3*grid_no-2,1);
j_idx = zeros(3*grid_no-2,1);
val   = zeros(3*grid_no-2,1);

i_idx(1:grid_no,1) = 1:grid_no;
j_idx(1:grid_no,1) = 1:grid_no;
val(1:grid_no,1)   = 4 * ones(grid_no,1);

i_idx(grid_no+1:2*grid_no-1,1) = 1:grid_no-1;
j_idx(grid_no+1:2*grid_no-1,1) = 2:grid_no;
val(grid_no+1:2*grid_no-1,1)   = -ones(grid_no-1,1);

i_idx(2*grid_no:3*grid_no-2,1) = 2:grid_no;
j_idx(2*grid_no:3*grid_no-2,1) = 1:grid_no-1;
val(2*grid_no:3*grid_no-2,1)   = -ones(grid_no-1,1);

m        = ( 3 * grid_no - 2 ) * grid_no + 2 * ( grid_no - 1 ) * grid_no;
m2       = 3 * grid_no - 2;
i_idx_A  = zeros(m,1);
j_idx_A  = zeros(m,1);
val_A    = zeros(m,1);

for ii = 1:grid_no
    i_idx_A((ii-1)*m2+1:ii*m2,1) = i_idx + ( ii - 1 ) * grid_no;
    j_idx_A((ii-1)*m2+1:ii*m2,1) = j_idx + ( ii - 1 ) * grid_no;
    val_A((ii-1)*m2+1:ii*m2,1)   = val;
end

i_idx_A(m2*grid_no+1:m2*grid_no+(grid_no-1)*grid_no,1) = 1:(grid_no-1)*grid_no;
j_idx_A(m2*grid_no+1:m2*grid_no+(grid_no-1)*grid_no,1) = grid_no+1:grid_no*grid_no;
val_A(m2*grid_no+1:m2*grid_no+(grid_no-1)*grid_no,1)   = -ones((grid_no-1)*grid_no,1);
i_idx_A(m2*grid_no+(grid_no-1)*grid_no+1:m,1)          = grid_no+1:grid_no*grid_no;
j_idx_A(m2*grid_no+(grid_no-1)*grid_no+1:m,1)          = 1:(grid_no-1)*grid_no;
val_A(m2*grid_no+(grid_no-1)*grid_no+1:m,1)            = -ones((grid_no-1)*grid_no,1);

mtx_A = sparse(i_idx_A, j_idx_A, val_A); 

%
% Compute right hand vector
%
rhs_f  = zeros(grid_no^2,1);
for jj = 1:grid_no
    for ii = 1:grid_no
        rhs_f((jj-1)*grid_no+ii) = 2 * pi^2 * sin(pi * x(ii)) * sin(pi * y(jj));
    end
end

rhs_f = leng_h^2 * rhs_f;

sol      = mtx_A \ rhs_f;
fprintf('CPU time for solving linear system = %14.4e, matrix size = %7.0f \n', toc, size(mtx_A,1));

s = RandStream('mt19937ar','Seed',0);
RandStream.setGlobalStream(s);
          
solver = 'SSOR'; %'GaussSeidel'; %'Jacobi';

switch solver
    case 'Jacobi'
        %
        n     = size( mtx_A, 1 );
        tol   = n * eps;
        maxit = 1500;
        it    = 1;
        flag  = 1;
        tmp   = rand(n,1);
        sol   = zeros(n,1);
        mtx_D = diag(diag(mtx_A));
        mtx_M = mtx_A - mtx_D;

        while ( it <= maxit && flag == 1 )
            sol      = mtx_D \ ( rhs_f - mtx_M * tmp ); 
            residual = norm(mtx_A * sol - rhs_f);
            fprintf('residual(%4.0f) = %17.10e \n', it, residual);
            if ( residual <= tol || norm(sol-tmp)/norm(sol) <= tol )
                flag = 0;
            else
                tmp  = sol;
                it   = it + 1;
            end
        end
    case 'GaussSeidel'
        %
        n     = size( mtx_A, 1 );
        tol   = n * eps;
        maxit = 1500;
        it    = 1;
        flag  = 1;
        tmp   = rand(n,1);
        sol   = zeros(n,1);

        while ( it <= maxit && flag == 1 )
            for ii = 1:n
                sol(ii,1) = rhs_f(ii) - mtx_A(ii,1:ii-1) * sol(1:ii-1,1) - mtx_A(ii,ii+1:n) * tmp(ii+1:n,1);
                sol(ii,1) = sol(ii,1) / mtx_A(ii,ii);
            end
            residual = norm(mtx_A * sol - rhs_f);
            fprintf('residual(%4.0f) = %17.10e \n', it, residual);
            if ( residual <= tol || norm(sol-tmp)/norm(sol) <= tol )
                flag = 0;
            else
                tmp  = sol;
                it   = it + 1;
            end
        end
    case 'SSOR'
        %
        omega   = 1.85;
        n       = size( mtx_A, 1 );
        tol     = n * eps;
        maxit   = 1500;
        it      = 1;
        flag    = 1;
        tmp     = rand(n,1);
        sol     = zeros(n,1);
        %
        % Take the strictly lower triangular matrix of matrix A
        %
        mtx_L   = tril(mtx_A, -1);
        %
        % Take the diagonal matrix of matrix A
        %
        diag_D  = diag(diag(mtx_A));
        %
        DomegaL = diag_D + omega * mtx_L;
        DL      = ( 1 - omega ) * diag_D - omega * mtx_L;

        %
        % Compute M(omega)^(-1) * b
        %
        M_inv_b = DomegaL' \ ( diag_D * ( DomegaL \ rhs_f ) );
        M_inv_b = ( omega * ( 2 - omega ) ) * M_inv_b;

        while ( it <= maxit && flag == 1 )
            %
            % Compute T(omega) * tmp
            %
            sol = DomegaL' \ ( DL * ( DomegaL \ ( DL' * tmp ) ) );
            %
            % Compute new vector sol + M_inv_b
            sol = sol + M_inv_b;
    
            residual = norm(mtx_A * sol - rhs_f);
            fprintf('residual(%4.0f) = %17.10e \n', it, residual);
            if ( residual <= tol || norm(sol-tmp)/norm(sol) <= tol )
                flag = 0;
            else
                tmp  = sol;
                it   = it + 1;
            end
        end
end


mtx_sol  = zeros(grid_no);
for ii = 1:grid_no
    mtx_sol(ii,:) = sol((ii-1)*grid_no+1:ii*grid_no);
end

figure(1)
mesh(xx, yy, zz);
xlabel('x');
ylabel('y');
zlabel('u');
title('u(x,y) = sin( x \pi ) sin( y \pi )');

figure(2)
surf(xx, yy, zz);
shading interp;
xlabel('x');
ylabel('y');
zlabel('u');
title('u(x,y) = sin( x \pi ) sin( y \pi )');

figure(3)
surf(xx, yy, mtx_sol);
shading interp;
xlabel('x');
ylabel('y');
zlabel('u');
title('Approximate solution');

figure(4)
surf(xx, yy, abs(mtx_sol-zz));
% shading interp;
xlabel('x');
ylabel('y');
zlabel('u');
title('Absolute Error');
