% Animation of time evolution of a superposition of the first three 
% energy eigenstates of the 1-D harmonic oscillator.  Relative weights of the 
% three eigenstates are set via parameters c0, c1, c2.

% Osciallator parameters
hbar=1;
m=1e3;
k=1;
omega=sqrt(k/m);
beta = sqrt(m*omega/hbar); 

% Weights in the superposition
c0 = 1/sqrt(3); 
c1 = 1/sqrt(3); 
c2 = 1/sqrt(3);

x = linspace(-10,10,1000);
% energies
V = k*x.^2/2; 
E = @(n) hbar*omega*(n+1/2); 
% wavefunctions
psi0 = sqrt(beta/sqrt(pi))*exp(-(beta*x).^2/2);
psi1 = sqrt(beta/(2*sqrt(pi)))*(2*beta*x).*exp(-(beta*x).^2/2);
psi2 = sqrt(beta/(8*sqrt(pi)))*(4*(beta*x).^2-2).*exp(-(beta*x).^2/2);
% probability density
rho = @(t) abs(c0*exp(-i*E(0)*t/hbar)*psi0 + c1*exp(-i*E(1)*t/hbar)*psi1 + ...
    c2*exp(-i*E(2)*t/hbar)*psi2).^2;

% animation
clear M;
offset = -1; 
periods = 3;
for t = 1:periods*2*pi/omega;
	plot(x,V,x,psi0+3*offset,x,psi1+2*offset,x,psi2+offset,x,rho(t));
    axis([-3,3,-4,4]);
	M(t) = getframe;
end
%v = VideoWriter('Harmonic012.avi');
%open(v);
%writeVideo(v,M);
%close(v);

%figure
%movie(M)

% E. Boyda / SMC Phys125 / Fall '17. 