% [u,x,t,kk,amp] = implicit_nonlinear_heat(t_0,t_f,dt,NPrint,N)
%
% this solves the semilinear heat equation 
%
%          u_t = u_xx + u^2
%
% with periodic boundary conditions using spectral derivatives in
% space and implicit time-stepping.  t_0 is the initial time, t_f is the
% final time, dt is the time-step size, N is the number of mesh-points, 
% and M is the number of printouts.  It also returns the power spectrum of the
% solution, where kk is the collection of wave numbers and amp = log(f_k^2)
% is the magnitude of the kth Fourier mode.
%
% N SHOULD BE A POWER OF 2!  
% (It will work no matter what, but will go much faster if N is a power
% of 2.)

function [u,x,t,kk,amp] = implicit_nonlinear_heat(t_0,t_f,dt,NPrint,N)

% define the mesh in space
dx = 2*pi/N;
x = 0:dx:2*pi-dx;
x = x';
% the wave-numbers...
kk = (-N/2+1):1:(N/2-1);

% total number of steps to be taken:
Ntot = (t_f-t_0)/dt;
% test to see if Ntot is an integer, to see if your choice of dt
% is compatible with your choice of t_0 and t_f
if abs(Ntot - round(Ntot)) > 100000*eps
  disp(' ')
  disp('dt is not consistent with your choice of t_0 and t_f.  Choose another dt!')
  disp(' ')
  break
end
Ntot = round(Ntot);
% number of steps to be taken before a (screen) printout
NRun = Ntot/NPrint;
% test to see if NRun is an integer to see if your choice of dt is
% consistent with NPrint
if abs(NRun - round(NRun)) > 1000*eps
  disp(' ')
  disp('NPrint is not consistent with dt.  Choose another NPrint!')
  disp(' ')
  break
end
NRun = round(NRun);
% define the vector t that is to be printed out...
t = t_0:NRun*dt:t_f;

% choose initial data
u_now = cos(3*x);

% We're computing solutions of
%
%    u_t = u_xx + |u|^2
%
% With implicit time-stepping, this would be
%
%  u_new - u_now = dt*u_new_xx + dt*u_now^2
%
%  u_new - dt*u_new_xx = u_now + dt*u_now^2
%
% That is, we want to solve
%
%      L u_new = RHS
%
% This is especially easily done spectrally since
%
%  (1 + dt k^2) u_new_hat(k) = RHS_hat(k)
%
% Hence
%
%  u_new_hat(k) = RHS_hat(k)/(1 + dt k^2) 

% save first printout:
u(:,1) = u_now;
v = fft(u_now); 
amp(N/2:N-1,1) = log10( abs(v(1:N/2)) + 1.e-16);
amp(1:N/2-1,1) = log10( abs(v(N/2+2:N)) + 1.e-16);

for jj=1:NPrint

  % take NRun time-steps
  for j=1:NRun

    % compute solution with one step of size dt
    u_coarse = tstep(u_now,dt,N);
    % compute solution with two steps of size dt/2
    u_fine = tstep(u_now,dt/2,N);
    u_fine = tstep(u_fine,dt/2,N);
    % use Richardson extrapolation to combine u_coarse and
    % u_fine to get a solution with local truncation
    % error of O(dt^3)...
    u_now = 2*u_fine - u_coarse;

  end

  % save printout
  u(:,jj+1) = u_now;
  v = fft(u_now); 
  amp(N/2:N-1,jj+1) = log10( abs(v(1:N/2)) + 1.e-16);
  amp(1:N/2-1,jj+1) = log10( abs(v(N/2+2:N)) + 1.e-16);
end
x(N+1) = 2*pi;
u(N+1,:) = u(1,:);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function u_now = tstep(u_now,dt,N);

    % compute u_now + dt u_now^2 
    wrk = u_now + dt*u_now.^2;
    % take its FFT and form the right-hand-side:
    RHS = fft(wrk);

    % we now time-step the PDE at the spectral level
    % by solving L v = RHS...
    ii = 1:N/2;
    k = (ii-1)';
    v(ii) = RHS(ii)./(1 + dt*k.^2);
    ii = N/2+2:N;
    k = (ii-N-1)';
    v(ii) = RHS(ii)./(1 + dt*k.^2);
    % take the inverse fourier transform.  
    u_now = real(ifft(v))';