%=====================================================================
% Solving PDE's in spectral space
%
% Computer Lab for Math 221, Fall 2002
% Sorin Mitran 09/24/2002
%=====================================================================

% Clear previous work
clc; clf; clear;

% I. Advection equation
%
%  q_t + u q_x = 0
%  q(x,t=0)=q0(x) 0<=x<=2 pi
%  u=1
n=64; dx=2*pi/n;
x=(0:n-1)*dx+dx/2;
t=0.; tfinal=2.0;
nmodesmax=4; u=1;

% Initial condition is built up as a sum of Fourier modes
for nmodes=1:nmodesmax
 q0=zeros(size(x)); a0=zeros([nmodes 1]); b0=zeros([nmodes 1]);
 for k=1:nmodes
   a0(k)=rand; b0(k)=rand;  
   q0 = q0 + a0(k) * cos(k*x) + b0(k) * sin(k*x);
 end 
 figure(1); plot(x,q0);
 in_ax=axis;
 txt = strcat('Initial condition. nmodes = ',num2str(nmodes));
 title(txt);
 xlabel('x'); ylabel('q0');
 % Fourier transform along x
 Qf0=fft(q0);
 figure(2); plot(1:n,abs(Qf0),'or',1:n,angle(Qf0),'xb');
 fin_ax=axis;
 title('Initial condition in Fourier space');
 xlabel('k'); ylabel('qhat');
 legend('abs(qhat)','arg(qhat)');
 clc; nmodes
 disp('Fourier modes included. Strike a key');
 pause
end

% Time loop
cfl=1;     % Change this to experiment with effect of different time steps
qf0=q0; qf1=qf0; q1=q0; qinit=q0; Qf0=fft(q0);
nt=1; dt=cfl*dx/u; sigma=u*dt/dx; I=sqrt(-1);
while t<tfinal
  % Plot solutions
  figure(1); plot(x,qinit,'.g',x,q1,'b',x+u*t,qinit,'o');
  axis(in_ax); title('Solution in real space');
  xlabel('x'); ylabel('q');
  legend('Initial condition','FTCS','Exact solution');
  % Compute using FT FS in real space
  for i=2:n
    q1(i) = q0(i) - sigma*(q0(i)-q0(i-1));
  end
  q1(1) = q0(1) - sigma*(q0(1)-q0(n));    
  %
  Qf1=fft(q1);
  figure(2); plot(1:n,abs(Qf1),'or',1:n,angle(Qf1),'xb')
  axis(fin_ax); title('Solution in Fourier space');
  xlabel('k'); ylabel('abs(qhat)');
  legend('abs(qhat)','arg(qhat)');
  clc; t=t+dt
  disp('time solution. Strike a key to continue ...');  
  pause  
  q0=q1;
end

disp('Strike a key to continue to heat equation example ...');  
pause  


% Heat equation
clc; clf; clear all;

% Problem parameters
fact=0.5;   % Change this to experiment with effect of different time steps
n=128; dx=2*pi/n;
x=(0:n-1)*dx+dx/2;
t=0.; tfinal=0.05; dt=fact*dx^2; sigma=dt/dx^2;
nmodes=20;
err=zeros([floor(tfinal/dt)+10 1]); eps=zeros(size(err));
% Initial condition
q0=zeros(size(x)); q1=zeros(size(x));
qf0=zeros(size(x)); qf1=zeros(size(x));
a0=zeros([nmodes 1]); b0=zeros([nmodes 1]);
for k=1:nmodes
  a0(k)=rand; b0(k)=rand;  
  q0 = q0 + a0(k) * cos(k*x) + b0(k) * sin(k*x);
end
qf0=q0;
figure(1); plot(x,q0);
xlabel('x'); ylabel('q');
txt = strcat('Initial condition for heat equation. nmodes=',num2str(nmodes));
title(txt);
in_ax=axis;
Qf0=fft(q0);
figure(2); plot(1:n,abs(Qf0),'or',1:n,angle(Qf0),'xb');
title('Initial condition Fourier coefficients')
xlabel('k'); ylabel('qhat');
legend('abs(qhat)','arg(qhat)');
fin_ax=axis;
disp('Initial condition. Strike a key');
pause
% Time loop
nt=1; Nsteps=floor(tfinal/dt)+2;
Qhist=zeros(nmodes,Nsteps);
while t<tfinal
  % Compute using FTCS
  for i=2:n-1
    q1(i) = q0(i) + sigma*(q0(i+1)-2*q0(i)+q0(i-1));
  end
  q1(1) = q0(1) + sigma*(q0(2)-2*q0(1)+q0(n));
  q1(n) = q0(n) + sigma*(q0(1)-2*q0(n)+q0(n-1));
  figure(1); plot(x,q1);
  axis(in_ax); hold on;  
  % Compute using Fourier x transform    
    % 1. Take x fft
    Qf0=fft(q0); mQ=zeros(size(Qf0));    % spectral time correction of FTCS initial condition
    mQexact=zeros(size(Qf0));
    % 2. Compute attenuation coefficient
    k2=n;
    for k1=2:n/2+1
      kmode2=(k1-1)^2;
      mQ(k1) = 1-dt*kmode2;
      mQ(k2) = 1-dt*kmode2;
      mQexact(k1)=exp(-dt*kmode2);
      mQexact(k2)=exp(-dt*kmode2);
      k2=k2-1;
    end    
    % 3. Attenuate Fourier coefficients
    Qf1 = Qf0.*mQ; Qf2 = Qf0.*mQ;
    Qhist(1:nmodes,nt)=abs(Qf1(1:nmodes))';
    % 4. Take x ifft
    qf1 = real(ifft(Qf1)); qf2 = real(ifft(Qf2)); 
  % Compare
  err(nt)=norm(qf1-q1); eps(nt)=err(nt)/norm(q1);
  figure(1); plot(x,qf1,'or',x,qf2,'.b');
  errtext = strcat('Relative err. = ',num2str(eps(nt)));
  xlabel('x'); ylabel('q'); title(errtext);
  legend('Real space FTCS','Fourier space FTCS','Fourier space exact');
  hold off; 
  figure(2); plot(1:n,abs(Qf1),'or',1:n,angle(Qf1),'xb')
  axis(fin_ax); title('Solution in Fourier space')
  xlabel('k'); ylabel('qhat');
  legend('abs(qhat)','arg(qhat)');
  % Iterate
  t=t+dt; nt=nt+1;
  q0=q1; qf0=qf1;
  %disp('Strike a key');
  pause(0.05);
end
figure(2);
plot(log10(eps),'xk');
xlabel('Iteration'); ylabel('eps')
title('Relative error');
figure(3)
surf(Qhist);
title('Fourier coefficient decay');
xlabel('t'); ylabel('k'); zlabel('abs(qhat)');