% Solves the heat equation u_t = u_xx with the Crank-Nicolson method
% The script should work in both Octave and MATLAB to the best of
% my knowledge.
% the Crank-Nicolson method is of order 2 in time and space.
% Anders Hoff - andershoff.net - 2010
% USE AT OWN RISK

clc;
clear all;
close all;

% predefinitions

% intervall a <= x <= b (space)
a = 0;
b = 1;

% number of steps in space
MM = 50 ;
M = MM - 1;
% stepsize in space
h = (b-a) / (M) ;
% number of steps in time
N = 50 ;
% time starts at (usually 0)
T1 = 0;
% max time
T2 = 0.2 ;
% stepsize in time
k = (T2 - T1) / (N-1) ;

% should not be necessary to change anything below
% remember to set the correct initial values and border values
% in the files f.m, g0.m and g1.m

e = ones(M-1,1) ;
% method is defined by the tridiagonal matrix A as well as the
% value in the vector G (from G.m) and the value r.
A = spdiags([ e -2*e e ] , -1:1 , M-1,M-1) ;
r = k / (2*h^2) ;

% internal x-values (space)
x = [a+h:h:b-h];
% all t-values (time)
t = [T1:k:T2] ;

% identity matrix (M-1) x (M-1)
I = eye(M-1);
% pre calculations
IpA = I + (r * A) ;
ImA = I - (r * A) ;

% assign initial distribution of heat. possble to simplify if x is a 
% proper function (or 0 as is the default setting)
for l= [1:1:M-1]
    u(l) = f(x(l)) ;
end
% column vector
u = u' ;
% U is the solution matrix
U = [g0(T1) u' g1(T2) ] ;

% solves M-1 equations for each N timesteps. (ie implicit method.)
for n = [1:1:N-1]
    u = ImA \ ( IpA*u + r .* G(t,M,n)  ) ;
    % assign current solution to end of solutions matrix U
    U = [U; [ g0(t(n)) u' g1(t(n)) ]] ;
end

% plot soltions
figure (1)
%mesh([a x b], t , U) ; % plot grid
surf([a x b], t , U) ; % plot solid mesh
xlabel('space');
ylabel('time');
zlabel('U')
title('Solutions of the heat equations with the Crank-Nicolson method.')

%END