%
% Quantum simple harmonic oscillator time evolver and plotter.  Same
% philosophy as the free particle and particle-in-a-box cases.  We
% basically write the initial wave function in terms of eigenfunctions
% and time evolve directly in the eigenbasis.
%

% basic constants set to one though we do track them correctly
clear
hbar=1;
omega=1;
sig=1;

% Size of box and x grid
L=30;
nx=400;
x=linspace(-L/2,L/2,nx)';
dx=x(2)-x(1);

% We only consider eigenstates up to n=nmax.  We then compute the Hermite
% polynomials ourselves as I was not able to find the function in matlab
% to give them simply.  hermiten has columsn indexing n and rows indexing
% the x values (really the dimensionless xi).  We know the first two
% Hermite polys (1 and 2x) and then use the recurrence relation to generate
% the rest.
nmax=100;
nvec=[0:nmax]';
hermiten = zeros(nx,nmax+1);
hermiten(:,1) = 1;
hermiten(:,2) = 2*x;
for n=1:nmax
    hermiten(:,n+2) = 2*x.*hermiten(:,n+1)-2*n.*hermiten(:,n);
end

% Energies of the eigenstates and their normalized wave functions.
E = hbar*omega*(nvec+0.5);
psin = zeros(nx,nmax+1);
for n=0:nmax
    psin(:,n+1) = hermiten(:,n+1).*exp(-x.^2/2);
    psin(:,n+1) = 1/sqrt(2^n*factorial(n)*sqrt(pi))*psin(:,n+1);
end

% Various choices for the intitial state
% (1) A Gaussian packet of width sig and momentum 7/sig:
%%psi0 = exp(-x.^2/sig^2).*exp(i*7/sig*x);
% (2) The ground state n=0 function displaced to x=7
%%psi0 = exp(-(x-7).^2 / 2);
% (3) The ground state n=0 function boosted by momentum hbar*7/sig
%%psi0 = exp(-x.^2 / 2).*exp(i*7/sig*x);
% (4) A coherent state with parameter labmda:  here labmda is real
% and equal to 5/sqrt(2) which means it is the same as displacing
% the n=0 ground state to x=5.
%%psi0 = zeros(size(x));
%%lambda=5/sqrt(2);
%%for n=0:nmax
%%psi0 = psi0 + exp(-lambda/2)*lambda^n/sqrt(factorial(n))*psin(:,n+1);
%%end

% normalize psi0 and find coefficients of eigenbasis expansion
psi0 = psi0/sqrt(sum(abs(psi0).^2)*dx);
c = psi0'*psin*dx;
c = c';

% Some values for plotting ranges
maxp = max(abs(psi0).^2)*1.1;
maxr = sqrt(maxp);

% We to nt time steps from t=0 to t=2*pi/omega which is the classical and
% quantum recurrence time (modulo a -1 phase overall coming from the
% zero-point energy: exp(-i*t*omega/2) at t=2*pi/omega is -1.
fs=20;
counter=0;
nt=80;
tmax=2*pi/omega;
for t=0:tmax/nt:tmax
    counter=counter+1;
    psit = psin*(c.*exp(-i*t*E/hbar));
    clf
    subplot(3,1,1)
    plot(x,real(psit),'b');
    set(gca,'fontsize',fs,'fontname','Arial')
    ylabel('Real \Psi')
    grid
    axis([-10 10 -maxr maxr]);
    title([sprintf('t = %.3f ',t/tmax) ' 2\pi/\omega'])
    subplot(3,1,2)
    plot(x,imag(psit),'r');
    set(gca,'fontsize',fs)
    ylabel('Imag \Psi')
    grid
    axis([-10 10 -maxr maxr]);
    subplot(3,1,3)
    plot(x,abs(psit).^2,'k');
    set(gca,'fontsize',fs,'fontname','Arial')
    xlabel('\xi','fontname','Symbol');
    ylabel('|\Psi|^2','fontname','Arial')
    grid
    axis([-10 10 0 maxp]);
    fname = sprintf('frame_%04d.jpg',counter-1);
    fprintf('counter=%04d/%4d  t/tmax=%.4f  %s\n',counter-1,nt,t/tmax,fname);
    eval(['print -djpeg -r200 ' fname]);
end