% finite element method for the heat equation with Dirichlet B.C.'s
% backward Euler in time

% set up u0
 Nexact = 10;
 x = chebfun('x',[0 1]);
 u0 = (x.^2).*(1-x);
 clear Psi      % eigenfunctions
 clear Phi      % hat functions

% find coefficients aexact_n(0) for exact solution
 aexact0 = zeros(Nexact,1);
 lambda = zeros(Nexact,1);
 for k=1:Nexact
    Psi{k} = sqrt(2)*sin(k*pi*x);
    lambda(k) = k^2*pi^2;
    aexact0(k) = (u0'*Psi{k})/(Psi{k}'*Psi{k});
 end

% set up finite element method
 Nfem = input('input N for finite element method: ');
 dt   = input('       time step for Euler method: ');   

 h = 1/(Nfem+1);
 xk = h*[1:Nfem]';
                             
% build stiffness matrix and mass matrix
 K = (1/h)*(2*eye(Nfem)-diag(ones(Nfem-1,1),-1)-diag(ones(Nfem-1,1),1));
 M = h*((2/3)*eye(Nfem)+(1/6)*diag(ones(Nfem-1,1),-1)+(1/6)*diag(ones(Nfem-1,1),1));
 A = -M\K;

% set up hat functions
 clear Phi
 for k=1:Nfem 
    Phi{k} = hat(k,Nfem);
 end  
 afem0 = u0(xk);

 tvec = [0:dt:.1];
 xx = [0;xk;1];                 % grid on which to plot the solution
 Ufem = zeros(length(xx),length(tvec));
 Uexact = zeros(length(xx),length(tvec));
 fprintf('press return to advance the solution in time\n')

 afem = afem0;                  % initial condition for Euler method
 I = eye(Nfem);
 for j=1:length(tvec)
    t = tvec(j);
    uexact = 0*x;
    ufem = 0*x;
    if j>1,                     % j=1 corresponds to t=0
       afem = (I-dt*A)\afem;    % backward Euler step
    end
    for k=1:Nexact
       uexact = uexact + aexact0(k)*exp(-lambda(k)*t)*Psi{k};
    end
    ufem = [0;afem;0];           % FEM solution at x grid pts = afem(t)

    figure(1), clf
    plot(uexact,'b-','linewidth',2)
    hold on
    plot(xx,ufem,'r-','linewidth',3)
    title(sprintf('exact and FEM+BE solution at $t=%10.5f$',t),...
       'fontsize', 18,'interpreter','latex')
    xlabel('$x$','fontsize',18,'interpreter','latex')
    ylabel('$u(x,t)$','fontsize',18,'interpreter','latex')
    Uexact(:,j) = uexact(xx);
    Ufem(:,j) = ufem;
    ylim([-.05 .25])
    pause
 end 

% waterfall plot of EXACT solution
 figure(2), clf
 plt=waterfall(xx,tvec,Uexact');
 set(plt,'edgecolor','k')
 xlabel('$x$','interpreter','latex','fontsize',18)
 ylabel('$t$','interpreter','latex','fontsize',18)
 zlabel('$u(x,t)$','interpreter','latex','fontsize',18)
 title('exact solution','interpreter','latex','fontsize',18)

% waterfall plot of FEM solution
 figure(3), clf
 plt=waterfall(xx,tvec,Ufem');
 set(plt,'edgecolor','k')
 xlabel('$x$','interpreter','latex','fontsize',18)
 ylabel('$t$','interpreter','latex','fontsize',18)
 zlabel('$u(x,t)$','interpreter','latex','fontsize',18)
 title('FEM+BE solution','interpreter','latex','fontsize',18)

% waterfall plot of error
 figure(4), clf
 plt=waterfall(xx,tvec,abs(Uexact-Ufem)');
 set(plt,'edgecolor','k')
 xlabel('$x$','interpreter','latex','fontsize',18)
 ylabel('$t$','interpreter','latex','fontsize',18)
 zlabel('$u(x,t)$','interpreter','latex','fontsize',18)
 title('error between exact and FEM+BE','interpreter','latex','fontsize',18)
