%  Compute theta-method solution to u_t = u_xx on [0,1] x [0,T]
%  with zero Dirichlet BC and pulse at x=1/2 as IC.

%------------------------------------------------------------------------
clear, close all
%
% enter parameters
%
T = 10;
asq = 1e4;  %used in IC (higher asq means a sharper spike)
N = 20;
%
theta = input('enter theta (0 <= theta <= 1): ');
%
if theta < 0 | theta > 1
  error('theta not in correct range')
elseif theta < 0.5
  rstar = 0.5/(1 - 2*theta);
else
  rstar = inf
end
%
r = input(['r* = ', num2str(rstar), '; enter r=dt*dx^2: ']);
%
%------------------------------------------------------------------------
% set other variables
%
dt = r/N^2;
ntmax = min(ceil(T/dt),1000);
%
s = 1 - 2*r;
M1 = - theta*r;
M0 = 1 + 2*theta*r;
S1 = (1 - theta)*r;
S0 = 1 - 2*(1-theta)*r;
%
x = [0:N].'/N;                    %% column vector
N1 = N + 1;
U = zeros(N1,ntmax+1);
%
%------------------------------------------------------------------------
% Set IC
%
uo = exp(-asq*(x-0.5).^2);
uo(1) = 0;
uo(N+1) = 0;                     %% force zero BC for IC
%
%========================================================================
% Calculate solution & store in vector U
%========================================================================
%
disp(' ')
disp(['Calculating solution up to time t = ',num2str(dt*ntmax),'...'])
U(:,1) = uo;
%
if theta == 0
%------------------------------------------------------------------------
% theta = 0: explicit scheme
%------------------------------------------------------------------------
  for n = 2:ntmax+1
    v = r*( uo(3:N1) + uo(1:N-1) ) + s*uo(2:N);
    uo(2:N) = v;
    U(:,n) = uo;
  end
else
%------------------------------------------------------------------------
% theta \in (0, 1]: implicit scheme
%------------------------------------------------------------------------
  e = ones(N-1,1);
  M = spdiags([M1*e M0*e M1*e], -1:1, N-1, N-1);
  S = spdiags([S1*e S0*e S1*e], -1:1, N-1, N-1);
  v = uo(2:N);
  for n = 2:ntmax+1
    q = S*v;
    v = M\q;
    U(2:N,n) = v;
  end
end
%
%========================================================================
% Make solution movie
%========================================================================
%
figure(1), clf,
set(1,'position',[407 12 612 684]);
set(1,'defaultlinelinewidth',3);
han =  plot(x,U(:,1),'b-');
xlabel('x', 'fontsize',20);
ylabel('U', 'fontsize',20);
title(['\theta = ', num2str(theta), '; r = ', num2str(r), '; N = ', int2str(N)], 'fontsize',20);
%
if r <= rstar
  iax = [-0.01, 1.0];
else
  iax = [-1,1];
end
axis([0,1, iax])
%
disp('Hit any key to play movie')
pause
%
set(han,'erasemode','xor')
for n = 2:ntmax+1
  set(han, 'ydata', U(:,n));
  pause(0.10)
  drawnow
end
