%
% Compute the finite difference approximation to the solution of the 2D advection 
% diffusion equation
%
%     - diff Delta y(x) + adv . nabla y(x) + react y(x) = f(x)  x in (0,1)^2,   
%                                                    y(x) = g(x)  x on the boundary
%          
% using matrix-free implementation of Jacobi and SOR using the stencil only.
%
% Matthias Heinkenschloss
% Jan 24, 2017   (revised Jan 22, 2020)
%



set(0, 'defaultaxesfontsize',24,'defaultaxeslinewidth',3.0,...
       'defaultlinelinewidth',3.0,'defaultpatchlinewidth',3.0,...
       'defaulttextfontsize',24);

clear all

%% Stopping tolerance for the iterative methods
tol = 1.e-4;


%% Problem data
diff   = 1.e-4;  % must be positive
theta  = 47.3 / 180 * pi;
adv    = [cos(theta),sin(theta)];  % must be component wise nonnegative
react  = 0;      % must be nonnegative

% set mesh points
n1  = 10;
n2  = 10;
h1  = 1/(n1+1);
h2  = 1/(n2+1);
x1  = (0:h1:1)';
x2  = (0:h2:1)';

%% Generate the 5-point stencil
% first generate the stencil as written in the notes
Stencil      = zeros(3,3);
Stencil(1,2) = -diff/h2^2;
Stencil(2,1) = -(diff+h1*adv(1))/h1^2;
Stencil(2,2) =  (2*diff+h1*adv(1))/h1^2 + (2*diff+h2*adv(2))/h2^2 + react;
Stencil(2,3) = -diff/h1^2;
Stencil(3,2) = -(diff+h2*adv(2))/h2^2;
% reorder so we can use Stencil.*y(i-1:i+1,j-1:j+1)
Stencil = [ Stencil(3,:)
            Stencil(2,:)
            Stencil(1,:) ];

% stencil using in Jacobi and SOR iteration
Stenciliter = Stencil; Stenciliter(2,2) = 0;

%% right hand side function  f(i,j) = right hand side function at (x_i, y_j)
f = zeros(n1+2,n2+2);
    
%% boundary data  g(i,j) = Dirichlet boundary data at (x1_i, x2_j)
g   = zeros(n1+2,n2+2);
g(:,1) = 1;    % g(x1,x2) = 1 for x2 = 0
ind = ( abs(x2) <= 0.25+eps );
g(1,ind) = 1;  % g(x1,x2) = 1 for x1 = 0 and x2 <= 0.25

%% unknown y(i,j) = unknown function at (x_i, y_j)
y = zeros(n1+2,n2+2);
y(1,:) = g(1,:); y(n1+2,:) = g(n1+2,:);  % set boundary values for unknown
y(:,1) = g(:,1); y(:,n2+2) = g(:,n2+2); 
y0     = y;

%% Jacobi method
res  = zeros(n1+2,n2+2);
y    = y0;
ynew = y;
for j = 2:n2+1
    for i = 2:n1+1
        res(i,j) = sum(sum(  Stencil.*y(i-1:i+1,j-1:j+1) ))  -  f(i,j);
    end
end
normres = norm(res(:),2);
reshist = normres;
while normres > tol
    % use lexicographic ordering
    for j = 2:n2+1
        for i = 2:n1+1
            ynew(i,j) = ( f(i,j) - sum(sum(  Stenciliter.*y(i-1:i+1,j-1:j+1) )) )/ Stencil(2,2);
        end
    end
    y = ynew;
    for j = 2:n2+1
        for i = 2:n1+1
            res(i,j) = sum(sum(  Stencil.*y(i-1:i+1,j-1:j+1) ))  -  f(i,j);
        end
    end
    normres = norm(res(:),2);
    reshist = [ reshist; normres ];
end

figure(1)
semilogy(reshist,'d-'); hold on
xlabel('Iteration')
ylabel('Residual Norm')


figure(2)
surf( x1, x2, y'); shading interp;
xlabel('x_1'); ylabel('x_2'); view(-60,45);
%print -depsc adv-diff-2d-ex1.eps




%% SOR method
for iomega = 1:2
    if( iomega == 1)
        fprintf(' use Gauss-Seidel \n')
        omega = 1;  % Gauss-Seidel first
    elseif( iomega == 2)
        fprintf(' use SOR \n')
        omega = 1.0014;  % use optimal SOR parameter from adv_diff_2d_sor.m
    end
    
    for order = 1:2
        
        % order == 1 use lexicographic ordering
        % order == 2 use red-black ordering
        if( order == 1)
            fprintf(' use lexicographic ordering \n')
        elseif( order == 2)
            fprintf(' use red-black ordering \n')
        end
        y     = y0;
        res   = zeros(n1+2,n2+2);
        for j = 2:n2+1
            for i = 2:n1+1
                res(i,j) = sum(sum(  Stencil.*y(i-1:i+1,j-1:j+1) ))  -  f(i,j);
            end
        end
        normres = norm(res(:),2);
        reshist = normres;
        while normres > tol
            if( order == 1)
                % use lexicographic ordering
                for j = 2:n2+1
                    for i = 2:n1+1
                        y(i,j) = (1-omega)*y(i,j) ...
                            + omega*( f(i,j) - sum(sum(  Stenciliter.*y(i-1:i+1,j-1:j+1) )) )/ Stencil(2,2);
                    end
                end
            elseif( order == 2)
                % use red-back ordering
                for j = 2:n2+1
                    for i = 2+mod(j,2):2:n1+1
                        y(i,j) = (1-omega)*y(i,j) ...
                            + omega*( f(i,j) - sum(sum(  Stenciliter.*y(i-1:i+1,j-1:j+1) )) )/ Stencil(2,2);
                    end
                end
                for j = 2:n2+1
                    for i = 2+(1-mod(j,2)):2:n1+1
                        y(i,j) = (1-omega)*y(i,j) ...
                            + omega*( f(i,j) - sum(sum(  Stenciliter.*y(i-1:i+1,j-1:j+1) )) )/ Stencil(2,2);
                    end
                end
            end
            for j = 2:n2+1
                for i = 2:n1+1
                    res(i,j) = sum(sum(  Stencil.*y(i-1:i+1,j-1:j+1) ))  -  f(i,j);
                end
            end
            normres = norm(res(:),2);
            reshist = [ reshist; normres ];
        end
        
        figure(1);
        if( iomega == 1 && order == 1)
            semilogy(reshist,'*-'); hold on;
        elseif( iomega == 1 && order == 2)
            semilogy(reshist,'+-'); hold on;
        elseif( iomega == 2 && order == 1)
            semilogy(reshist,'o--'); hold on;
        elseif( iomega == 2 && order == 2)
            semilogy(reshist,'square--'); hold on;
        end
        xlabel('Iteration')
        ylabel('Residual Norm')
        
        figure(2)
        surf( x1, x2, y'); shading interp;
        xlabel('x_1'); ylabel('x_2'); view(-60,45);
        %print -depsc adv-diff-2d-ex1.eps
        
    end % order
end % iomega
figure(1); hold off;
legend('Jacobi','GS-lex','GS-rb','SOR-lex','SOR-rb', ...
       'Location','NorthEast')



