% M-file: wave_packet_cn.m 
% Evolves a Gaussian wave packet according to the Schrodinger equation in 1D 
% using the Crank-Nicolson numerical algorithm


% PARAMETERS:
% IMPORTANT: Matlab assumes that i is the IMAGINARY UNIT i=sqrt(-1), so do not
% use it for anything else!  

% The space grid along the x-axis has K_space points: 
  K_space=2500;

% The time grid has N_time steps:
  N_time=2048;

% Space-Time grid spacings:
  dx=5.0e-4; dt=5.0e-7; 

% Wave packet parameters:
  sigma=0.016; k0=700.0;  x0=0.3;
  
% Barrier parameteres:
  bar_thick=0.1; % thickness
  V0=2.0*k0*k0; % height
  

% VARIABLES:
% Space-Time indices:
  k=0; n=0;

% TDSE variable:
  q=0;

% ARRAYS:
% Space grid x(i) and potential V(x) evaluated this space grid
  x=zeros(K_space,1); V=zeros(K_space,1);

% Wave function Psi(x,t) evaluated on discrete space-time grid is stored as complex array:
  Psi=zeros(K_space,N_time);

  % Probability density defined by |Psi|^2:
  rho=zeros(K_space,1);

% Auxilary complex matrices of discrete Time-Dependent Schrodinger Equation (TDSE):
  A=zeros(K_space,K_space); B=zeros(K_space,K_space);


% TDSE parameter:
  q=dt/(4*dx*dx);

  
% Evaluate coordinates of the discrete space grid:
  x(1)=0;
  for k=1:K_space-1
    x(k+1)=x(k)+dx;
  end

% Evaluate potential V(x) on the space grid x(k) and load into array V(1:K_space)
  for k=1:K_space
          
    if ((x(k)>0.6) && (x(k)<0.6+bar_thick))
      V(k)=V0;
    end

  end


% Load initial Gaussian wave packet at time t=0 (which corresponds to n=1) into Psi(:,1)

  Psi(:,1)=(1.0/((2.0*pi)^0.25*sigma^0.5))*exp((-(x-x0).^2)/(4*sigma^2)).*exp(i*k0*(x-x0));

% Note this is MATLAB shortcut for a loop which in other programming languages will be 
% have to be implemenet as, e.g.,:  
% for k=1:K_space
%     Psi(k,1)=(1.0/((2.0*pi)^0.25*sigma^0.5))*exp((-(x(k)-x0)^2)/(4*sigma^2))*exp(i*k0*(x(k)-x0));
% end;

% Here you should set up  matrices A and B following TDSE expression:

% A(1,1)=?;
% B(1,1)=?;

% for k=2:K_space
%    A(k,k)=?;
%    A(k-1,k)=?;
%    A(k,k-1)=?;
%    B(k,k)=?;
%    B(k-1,k)=?;
%    B(k,k-1)=?;
% end



% Solve TDSE A*Psi(n+1)=B*Psi(n), using LU decomposition, for N_time-1 steps starting from the minimum 
% uncertainty (Gaussian wave packet) state loaded above at t=0, time step n=1:

  for n=1:N_time-1
     Psi(:,n+1)=A\(B*Psi(:,n));
  end


% Select wave function at time step n and plot the probability density defined by it: 
n=1;
rho=conj(Psi(:,n)).*Psi(:,n);

[AX,H1,H2] = plotyy(x,rho,x,V,'plot');
set(get(AX(1),'Ylabel'),'String','Probability Density'); 
set(get(AX(2),'Ylabel'),'String','Potential');
xlabel('X-axis coordinate');
title(strcat('Wave packet at time step=',num2str(n))); 


% You can also make a Matlab movie of the wave packet propagation:
axis tight
set(gca,'nextplot','replacechildren')
disp('Creating the movie')
for n=1:N_time-1
    rho=conj(Psi(:,n)).*Psi(:,n);
    plot(x,rho)
    F(n)=getframe;
end
disp('Playing the movie')
movie(F)