% DIFFUSION: Simulate 2D diffusion equation.

% Numerical scheme used is a first order upwind in time and a second
% order central difference in space (Implicit and Explicit)

% Suraj Shankar, 2012

%% Parameters

nx = 40;                           %Number of steps in space(x)
ny = 50;                           %Number of steps in space(y)
nt = 30;                           %Number of time steps

dt = 0.01;                         %Width of each time step
dx = 2/(nx-1);                     %Width of space step(x)
dy = 2/(ny-1);                     %Width of space step(y)

x = 0:dx:2;                        %Range of x(0,2) and specifying the grid points
y = 0:dy:2;                        %Range of y(0,2) and specifying the grid points

u  = zeros(nx,ny);                 %Preallocating u
un = zeros(nx,ny);                 %Preallocating un

vis = 0.1;                         %viscosity (diffusion coefficient)

UW = 0;                            %x=0 boundary condition
UE = 0;                            %x=L boundary condition
US = 0;                            %y=0 boundary condition
UN = 0;                            %y=L boundary condition


%% Initial conditions (in time)

% Start with u = 2 in a region that's 1/2-by-1/2, u = 0 elsewhere.

for i = 1:nx
    for j = 1:ny
        if ((1<=y(j))&&(y(j)<=1.5)&&(1<=x(i))&&(x(i)<=1.5))
            u(i,j) = 2;
        else
            u(i,j) = 0;
        end
    end
end

%% Boundary conditions (in space)

% boundary condition vector
bc = zeros(nx-2,ny-2);
bc(1,:) = UW/dx^2; 
bc(nx-2,:) = UE/dx^2;
bc(:,1) = US/dy^2; 
bc(:,ny-2) = UN/dy^2;

% boundary conditions at the corners
bc(1,1) = UW/dx^2 + US/dy^2; 
bc(nx-2,1) = UE/dx^2 + US/dy^2;
bc(1,ny-2) = UW/dx^2 + UN/dy^2; 
bc(nx-2,ny-2) = UE/dx^2 + UN/dy^2;

% scale boundary conditions for viscosity
bc = vis*dt*bc;

% coefficient matrix for the implicit scheme
Ex = sparse(2:nx-2, 1:nx-3, 1, nx-2, nx-2); % 1's below diagonal
Ax = Ex + Ex' - 2*speye(nx-2);              % 1 -2 1 tridiagonal

Ey = sparse(2:ny-2, 1:ny-3, 1, ny-2, ny-2);
Ay = Ey + Ey' - 2*speye(ny-2);

A = kron(Ay/dy^2, speye(nx-2)) + ...  % build 2D matrix from 1D matrices
    kron(speye(ny-2),  Ax/dx^2);

I = speye((nx-2)*(ny-2));
D = I - vis*dt * A; 

%% Main loop : stepping in time

i = 2:nx-1;    % indices of interior points
j = 2:ny-1;    % indices of interior points

for it = 0:nt
    
    un = u;
    
    % plot the picture at this time step
    h=surf(x,y,u','EdgeColor','none');      
    shading interp
    axis ([0 2 0 2 0 2])
    title({['2-D Diffusion with {\nu} = ',num2str(vis)];['time (\itt) = ',num2str(it*dt)]})
    xlabel('Spatial co-ordinate (x) \rightarrow')
    ylabel('{\leftarrow} Spatial co-ordinate (y)')
    zlabel('Transport property profile (u) \rightarrow')
    drawnow;
    refreshdata(h);
    pause(.1);
    
    if implicit_method
        
        U = un; 
        U(1,:) = []; U(end,:) = []; % trim off the edges of U
        U(:,1) = []; U(:,end) = []; % trim off the edges of U
        U = U(:);                   % reshape U as a vector
        U = D\U;                    % All the work is right here!
        u(2:nx-1,2:ny-1) = reshape(U, nx-2, ny-2);
        
    else % explicit method
        
        u(i,j) = un(i,j) + ...
            (vis*dt * (un(i+1,j) -2*un(i,j) +un(i-1,j)) / (dx*dx) ) + ...
            (vis*dt * (un(i,j+1) -2*un(i,j) +un(i,j-1)) / (dy*dy) );
    end; 
    
    % enforce boundary conditions
    u(1,:)  = UW;
    u(nx,:) = UE;
    u(:,1)  = US;
    u(:,ny) = UN;
    
end