function [u,tau]= cnheat(ndiv,nstep,r)
% Crank-Nicolson solution of 1-D heat equation
%   input parameters
%          ndiv      # of spatial subdivisions
%          nstep     # of time steps
%          r         mesh ratio
%   output parameters
%          u         solution vector
%          tau       final time
h=1/ndiv; tau=0; x=h:h:1-h;
fprintf('subdivision parameter : %g\n',h)
% define and plot initial condition
u=min(2*x,2-2*x); stepsolplot(tau,x,u)
n=length(u);  k=r*h*h;
fprintf('timestep              : %g\n',k)
% timestep loop
for step=1:nstep
    tau=tau+k;
    alpha=1-r; beta=r/2;
    v(1)=alpha*u(1)+beta*u(2);
    for i=2:n-1
        v(i)=beta*u(i-1)+alpha*u(i)+beta*u(i+1);
        end
    v(n)=beta*u(n-1)+alpha*u(n);
    gamma=-r/2; alpha=1+r;beta=-r/2;
    u=trisolve(gamma,alpha,beta,v);
    stepsolplot(tau,x,u)
    end

function stepsolplot(tau,x,u)
if tau==0, % plot initial condition
   plot([0,x,1],[0,u,0],'-xk'), axis('square'), title('initial profile'),
   xlabel('x'), ylabel('temperature'), shg, pause(3)
else % plot solution
   hold on, plot([0,x,1],[0,u',0],'-ob'),
   title(['Crank-Nicolson: time is ',num2str(tau)])
   pause(0.2), hold off
end
return

function x = trisolve(gamma,alpha,beta,b)
%trisolve  solves tridiagonal system Ax=b
%   input parameters
%          gamma     subdiagonal coefficient
%          alpha     diagonal coefficient
%          beta      superdiagonal coefficient
%          b         right hand side vector
n=length(b); d=ones(n,1)*alpha; x=zeros(n,1);
for k = 1:n-1
   pivot=d(k);
   mult=gamma/pivot;
   d(k+1)=d(k+1)-mult*beta;
   b(k+1)=b(k+1)-mult*b(k);
   end
x(n)=b(n)/d(n);
for k = n-1:-1:1
   x(k)=(b(k)-beta*x(k+1))/d(k);
   end
return
                  