%% Program for solving u_{t} = Du_{xx} on x in [a,b] for 0<=t<=t_end.
clear all
close all
clc
D = 2;

a = 0;
b = 1;
h = 0.1;

t0 = 0;
t_end = 1;
k = 0.002;

M = (b-a)/h;
m = M - 1;
N = t_end / k;

% u0 = @(x) 2 * cosh(x);
% ua = @(t) 2*exp(2*t);
% ub = @(t) (exp(2) + 1) * exp(2*t - 1);

u0 = @(x) exp(x);
ua = @(t) exp(2*t);
ub = @(t) exp(2*t + 1);

% u_sol = @(x,t) exp(2*t+x)+exp(2*t-x);
u_sol = @(x,t) exp(2*t+x);
T = t0 + (0:N)*k;
x = a + (1:m) * h;
X = [a x b];


sigma = D * k / h^2;
disp(['Sigma = ' num2str(sigma)])
A = diag(1 - 2*sigma*ones(m,1)) + diag(sigma*ones(m-1,1),1)+diag(sigma*ones(m-1,1),-1);




figure(1)
U = zeros(m,N+1);
U(:,1) = u0(x)';
for j = 1:N
    U(:,j + 1) = A * U(:,j) + sigma * [ua(T(j));zeros(m-2,1);ub(T(j))];
end


U_mesh = [ua(T);U;ub(T)];
[T_mesh,X_mesh] = meshgrid(T,X);

mesh(T_mesh,X_mesh,U_mesh)
xlabel('T')
ylabel('X')
zlabel('U')
hold on
plot3(zeros(M+1),X,u0(X),'r')
plot3(ones(M+1),X,u_sol(X,1),'r')
hold off
figure(2)
plot(X,abs(u_sol(X,1)'-U_mesh(:,end)),'r')

