%{
 Solving the diffusion equation ft = a*fxx
 Methods: 
 Backward-time Centered-space difference
 AG: April 2022
%}

clear  % Clears command window
clc    % Clears command history
clf    % Removes anything in the figure window before simulation. 
tic    % set timer for the program

% INPUT
% choose method key =
% 1 - FTCS, 2 - BTCS
key = 2;
% Physics
a = 0.01;  % diffusion coefficient
% the domain of solutions
xmin = 0.0;
xmax = 1.0;
tmin = 0.0;
tmax = 10.;
% the grid
nx =  21;
nt = 101;
% 
x = zeros(1,nx);
y = zeros(1,nt);
f = zeros(nx,nt);
dx = (xmax-xmin)/(nx-1);
dt = (tmax-tmin)/(nt-1);
% INITIAL and BOUNDARY conditions
% initial condition f(x,0)
for i = 1 : nx
  x(i) = xmin + dx*(i-1);
  f(i,1) =  sin(pi*x(i)) + x(i);
end
% Boundary conditions f(0,t) and f(L,t)
for j = 1 : nt
  y(j) = tmin + dt*(j-1);
  f(1,j) =  0.0;
  f(nx,j) = 0.0;
end

figure (1)
surf(y,x,f)
xlabel('Time'), ylabel('Space'), zlabel('Solution')

% call the solver =================
if key == 1        % FTCS method
  [f2] = pdeP1(f,dx,dt,nx,nt,a);
elseif key == 2    %BTCS method
  [f2] = pdeP2(f,dx,dt,nx,nt,a);
end

figure (2)
surf(y,x,f2)
xlabel('Time'), ylabel('Space'), zlabel('Solution')

%print time of calculations
timeE = floor(toc);
fprintf('\n timeElapsed = %4i sec \n',timeE)

%{
  Solving the diffusion equation with Dirichlet BCs
  Method: Forward-time Centered-space difference
  INPUT:
  f(i,j)    boundary conditions 
  dx, dy    grid increments
  nx        number of grid points in x direction 
  nt        number of grid points in y direction
  a         diffusion coefficient
  OUTPUT
  f2(x,y)   the solution
%}
function[f2] = pdeP1(f,dx,dt,nx,nt,a)
% preparation
f2 = zeros(nx,nt);
d = a*dt/(dx^2);
fprintf(' Diffusion parameter d = %6.4f \n',d)
if d > 0.5
  fprintf(' ATTENTION: the diffusion parameter is too large for FTCS method \n ')
end

for j=1:nt
  for i=2:nx-1
    f(i,j+1) = f(i,j)+d*(f(i+1,j)-2.0*f(i,j)+f(i-1,j));
  end
end
% Prepare OUTPUT 
for j=1:nt
  for i=1:nx
    f2(i,j)= f(i,j);
  end
end
end


%{
  Solving the diffusion equation with Dirichlet BCs
  Method: Backward-time Centered-space difference
  INPUT:
  f(i,j)    initial matrix with boundary conditions 
  dx, dy    grid increments
  nx        number of grid points in x direction 
  nt        number of grid points in y direction
  alpha     diffusion coefficient
  OUTPUT
  f2(x,y)   the solution
AG: April 2022
%}
function[f2] = pdeP2(f,dx,dt,nx,nt,alpha)
% preparation
a  = zeros(nx-2,3);
b  = zeros(1,nx-2);
f2 = zeros(nx,nt);
d = alpha*dt/(dx^2);
fprintf(' Diffusion parameter d = %6.4f \n',d)

for j=2:nt
  for i=2:nx-1
    a(i-1,1) = 0.0-d;
    a(i-1,2) = 1.0+2.0*d;
    a(i-1,3) = 0.0-d;
    b(i-1) = f(i,j-1);
  end
    a(1,1) = 0.0;
    b(1) = b(1) + d*f(1,j);
    a(nx-2,3)= 0.0;
    b(nx-2) = b(nx-2) + d*f(nx,j);

  [w] = Thomas(a,b,nx-2);
  for i=1:nx-2
    f(i+1,j) = w(i);
  end
end
% Prepare OUTPUT 
for j=1:nt
  for i=1:nx
    f2(i,j)= f(i,j);
  end
end
end

function [x] = Thomas(c,b,n)
% !============================================================
% ! Solutions to a system of tridiagonal linear equations c*x=b
% ! Method: the Thomas method
% ! Alex G. March 2021
% !-----------------------------------------------------------
% ! input ...
% ! c(n,3) - array of coefficients for matrix C
% ! b(n)   - vector of the right hand coefficients b
% ! n      - number of equations
% ! output ...
% ! x(n)   - solutions
% ! comments ...
% ! the original arrays c(n,3) and b(n) will be destroyed 
% ! during the calculation
% !===========================================================

%step 1: forward elimination
for k = 2:n
  coeff=c(k,1)/c(k-1,2);
  c(k,2)=c(k,2)-coeff*c(k-1,3);
  b(k)=b(k)-coeff*b(k-1);
end

%step 2: back substitution
x(n) = b(n)/c(n,2);
for k = n-1:-1:1
    x(k) = (b(k)- c(k,3)*x(k+1))/c(k,2);
end
end






%  f(1,1)=0;   f(2,1)=20; f(3,1)=40; f(4,1)=60; f(5,1)=80;
%  f(6,1)=100; f(7,1)=80; f(8,1)=60; f(9,1)=40; f(10,1)=20; f(11,1)=0;