%
% Particle in a box simulator: given an initial psi(x,0)=psi0, it
% evolves it in time and for each time slice plots Real psi(x,t), Imag,
% and |Psi|^2.  We plot quantities in terms of scaled variables (e.g. x/a).
% 
% The box is 0<x<a.
% The initial wave function is psi0.
% Normalizations are done with the simplest possible integration scheme:
%   Integral { f(x) dx } = sum_i f(x_i)*dx
% where x_i are evenly spaced from 0 to a and dx is just the uniform
% grid spacing.
%
% The actual time evolution is done in the eigenbasis:  we first
% write psi0 in terms of the particle-in-a-box eigenfunctions and
% find the expansion coefficients c(:).  Then we just use the
% time evolution formula with phase factor exp(-i*t*En/hbar) to sum
% them back up.  The code stores many things as full matrices: thus
% it should be fast but memory intensive.
%

% Just set fundamental values to unity for simplicity although we
% are careful to work with proper dimensions below in the formulae.
clear
a=1;
hbar=1;
m=1;

% First energy level of particle-in-a-box and the first recurrence time
E1=hbar^2*pi^2*(1^2)/(2*m*a^2);
trec = 2*pi*hbar/E1;

% We will only use the first nmax eigenstates and assume c(n>nmax)=0.
% The higher nmax, the higher the fidelity of the calculation.  E(:)
% contains the energies of the states En
nmax=400;
E = E1*[1:nmax].^2;
E = E';

% Here we create the x axis by making a uniform grid of x values from 0
% to a with nx points.  We then get spacing and use the formula
% for the eigenstates. psin holds the eigenstates as a matrix:  the columns
% are indexing the states n and the rows are the x values.
nx=500;
x=linspace(0,a,nx)';
dx=x(2)-x(1);
nvec=[1:nmax]';
psin = sqrt(2/a)*sin(pi*x*nvec'/a);

% initial wave function
sig=0.15;
k0=50;
psi0 = exp(-((x-0.5)/sig).^2).*exp(i*k0*x);
c = psi0'*psin*dx;
c = c';
E0 = sum(abs(c).^2 .* E);
fprintf('Initial packet has energy E0 = %g with sigE = %g\n',...
    E0,sqrt(sum(abs(c).^2 .* E.^2)-E0^2));

% Data on maximum |Psi(x,0)|^2 so we can set the plot bounds

maxp = max(abs(psi0).^2);
maxr = sqrt(maxp);

% We will be doing nt time steps from zero to the recurrence time trec.
fs=20;
counter=0;
nt=1000;
for t=0:trec/nt:trec
    counter=counter+1;
    % compute the psi(x,t) into psit
    psit = psin*(c.*exp(-i*t*E/hbar));
    % make plots stacked vertically: real, imag, and square
    clf
    subplot(3,1,1)
    plot(x,real(psit),'b');
    set(gca,'fontsize',fs)
    ylabel('Real \Psi')
    grid
    axis([x(1) x(end) -maxr maxr]);
    title([sprintf('t = %.3f',t/trec') ' * h/E_1'])
    subplot(3,1,2)
    plot(x,imag(psit),'r');
    set(gca,'fontsize',fs)
    ylabel('Imag \Psi')
    grid
    axis([x(1) x(end) -maxr maxr]);
    subplot(3,1,3)
    plot(x,abs(psit).^2,'k');
    set(gca,'fontsize',fs)
    xlabel('x/a');
    ylabel('|\Psi|^2')
    grid
    axis([x(1) x(end) 0 maxp]);
    % write out to a jpeg file by first creating a name with the
    % "frame" (i.e. time) number and then printing to it.  Other
    % approaches would be to use the getframe function and either
    % play the movie at the end with the movie function or use
    % imwrite to write the images individually.  The print approach
    % is slower but more robust.
    fname = sprintf('frame_%04d.jpg',counter-1);
    eval(['print -djpeg -r150 ' fname]);
    % screen output to see progress
    fprintf('counter=%04d/%4d  t/trec=%.4f  %s\n',counter,nt,t/trec,fname);
end