% Lord_Kelvin_CN.m
%
% Solves the 1D heat equation with Crank-Nicholson finite difference scheme

clear

% Physical parameters
L       =   1e6;		%   length of modeled domain     [m]     
Tini    =   2300;		%   initial Earth temperature    [K]
Tsurf   =   300;		%   surface temperature          [K]
kappa   =   1e-6;		%   thermal diffusivity of rocks [m2/s]
day     =   3600*24;		%   # seconds per day
year    =   day*365;		%   # days per year
maxtime =   1e8*year;		%   max. time to run             [s]

% Numerical parameters
nz      =   101;		%   number of gridpoints in z-direction
nt      =   200;		%   number of timesteps to compute
dt      =   maxtime/nt;		%   timestep			 [s]
dz      =   L/(nz-1);		%   grid spacing "h"
z       =   0:dz:L;		%   grid

% Setup initial temperature profile
T       =   ones(size(z))*Tini;
Tnew    =   zeros(1,nz);	% alternative way to initialize as vector
T(1)    =   Tsurf;		% Dirichlet boundary condition on top
T(nz)   =   Tini;		% Dirichlet boundary condition deep mantle

% Setup time array to plot the thermal gradient
ti      =   NaN(1,nt);		% 'NaN' instead of 'zeros' to avoid misplot
Tgrad   =   zeros(1,nt);	% numerical solution
Tana    =   zeros(1,nt);	% analytical solution
done    =   0;			% for estimation of Earth age

% Setup propagation constant
beta = kappa*dt/(dz*dz);

% Setup lhs matrix A for BTCS
% A = sparse(nz,nz);
% for i=2:nz-1
%     A(i,i-1) = - beta;
%     A(i,i)   = (1 + 2*beta);
%     A(i,i+1) = - beta;
% end
% A(1,1)   = 1;
% A(nz,nz) = 1;

% Setup rhs matrix B for BTCS
% B = sparse(nz,nz);
% for i=2:nz-1
%    B(i,i-1) = 0;
%    B(i,i)   = 1;
%    B(i,i+1) = 0;
% end
% B(1,1)   = 1;
% B(nz,nz) = 1;

% Setup rhs vector for BTCS
% rhs     = zeros(nz,1);
% rhs(1)  = Tsurf;
% rhs(nz) = Tini;
% rhs     = B*rhs;

% Setup lhs matrix A for Crank-Nicholson
A = sparse(nz,nz);
for i=2:nz-1
    A(i,i-1) = - beta;
    A(i,i)   = 2*(1+beta);
    A(i,i+1) = - beta;
end
A(1,1)   = 1;			% needs to be checked !
A(nz,nz) = 1;			% needs to be checked !

% Setup rhs matrix B for Crank-Nicholson
B = sparse(nz,nz);
for i=2:nz-1
   B(i,i-1) = beta;
   B(i,i)   = 2*(1-beta);
   B(i,i+1) = beta;
end
B(1,1)   = 1;			% needs to be checked !
B(nz,nz) = 1;			% needs to be checked !

% Setup rhs vector for Crank-Nicholson
rhs     = zeros(nz,1);
rhs(1)  = Tsurf;
rhs(nz) = Tini;
rhs     = B*rhs;


% loop over time
time    =   0;
for n=1:nt			% Timestep loop

    % FTCS scheme
    % Tnew(2:(end-1)) = T(2:(end-1)) + ... 
    %  beta*(T(1:(end-2))+T(3:end)-2*T(2:(end-1)));

    % BTCS scheme
    % rhs(2:nz-1) = T(2:nz-1);
    % Tnew = A\rhs;		% Gauss elimination

    % Crank-Nicholson scheme
    rhs(2:nz-1) = T(2:nz-1);
    rhs = B*rhs;
    Tnew = A\rhs;		% Gauss elimination

    % Set boundary conditions
    Tnew(1)     =   T(1);
    Tnew(nz)    =   T(nz);
    
    % Update temperature and time
    T           =   Tnew;
    time        =   time+dt;
        
    % compute thermal gradient in K/km
    ti(n)       =   time;
    Tgrad(n)    =   1000*(T(2)-T(1))/dz;
    Tana(n)     =   1000*(Tini-Tsurf)/sqrt(pi*kappa*time);
    
    % print some progress information
    % if (n == nt)
    % fprintf('%g out of %g: Tgrad = %f  Tana = %f  T(2) = %f\n', ...
    %          n,nt,Tgrad(n),Tana(n),T(2));
    % end
    
    % evaluate the Earth age for 25 K km^{-1}
    check = Tgrad(n)-25;
    if (done == 0)
       if (check < 0)
          fprintf('estimated Earth age = %g Mill. years\n\n',time/(1e6*year));
          done = 1;
       end
    end
    
    % Plot solution
    figure(1), clf
    subplot(1,2,1)
    plot(T,-z/1000)
    grid
    xlabel('temperature [^oK]')
    ylabel('z [km]')
    title(['temperature T(z)'])
    subplot(1,2,2)
    plot(ti/(1e6*year),Tgrad,ti/(1e6*year),Tana)
    axis([0 100 0 100])
    xlabel('time [Myr]')
    ylabel('\DeltaT/\Deltaz_{num} [K km^{-1}]')
    title(['geothermal gradient'])
    
    drawnow
end
