% @(#) 6.728 Homework Problem 3.3 -- Matlab Code
% ps3p3.m file
% @(#) This program computes decomposes an the initial state of a particle
% in a box state into the M lowest energy eigenstates (i.e., cosines and
% sines). It then performs a time evolution.
% @(#) See Also: Eigenfunction Expansion Tutorial

% preamble
diary ps3p3.out % keep a log of our output
diary on;

clear all; % Clear workspace
clc; % Clear the command window
close all; % Get rid of old figures
format compact; % suppresses extra lines

% Defines some standard constants used in quantum mechanics and
% statistical mechanics.
h = 6.62607E-34; % Planck constant [J*s]
hbar = h/(2*pi); % Rationalized Planck constant [J*s]
Me = 9.1094E-31; % Electron rest mass [kg]
Mp = 1.672623E-27; % Proton rest mass [kg]
Mn = 1.674929E-27; % Neutron rest mass [kg]
c = 2.99792458E8; % Speed of light in vacuum [m/s]

L = 1; % 1D Box [-1/2, 1/2] is of unit length
n_max = 20;

% Create our x vector
dx= 0.5E-3;
x = [-0.5:dx:0.5]';
Lx = length(x);

% Create the column vectors psi.

psi = zeros(Lx,1);
b_left = max(find(abs(x+1/16) < dx));
b_right = min(find(abs(x-1/16) < dx));
psi(b_left:b_right) = sqrt(8);
% Initial state: psi(x,0) = \sqrt{\frac{8}{L}} if |x| \leq L/16

psi_pdf = conj(psi).*psi;

% Create matrix of basis functions phi where energy level n is incremented
% column by column and location x is incremented row by row.
%
% phi_1(x(1)) phi_2(x(1)) ... phi_{n_max}(x(1))
% phi = phi_1(x(2)) phi_2(x(2)) ... phi_{n_max}(x(2))
% | | \ |
% phi_1(x(Lx)) phi_2(x(Lx)) ... phi_{n_max}(x(Lx))
%
% in our case
% phi_{n odd} = sqrt{2/L) * cos(n*pi*x/L)
% phi_{n even} = sqrt(2/L) * sin(n*pi*x/L)
%
phi = zeros(Lx,n_max); % allocate space for phi matrix
N = ones(Lx,1)*[1:n_max];
X = x*ones(1,n_max);

if rem(n_max,2) == 0 % that is, if n_max is even
phi(:,1:2:n_max-1) = sqrt(2/L)*cos(N(:,1:2:n_max-1).*pi.*X(:,1:2:n_max-1)/L); % odd columns are cosines
phi(:,2:2:n_max) = sqrt(2/L)*sin(N(:,2:2:n_max).*pi.*X(:,2:2:n_max)/L); % even columns are sines
else % that is, if n_max is odd
phi(:,1:2:n_max) = sqrt(2/L)*cos(N(:,1:2:n_max).*pi.*X(:,1:2:n_max)/L); % Again, odd columns are cos
phi(:,2:2:n_max-1) = sqrt(2/L)*sin(N(:,2:2:n_max-1).*pi.*X(:,2:2:n_max-1).L); % even columns are sines
end

a = dx * phi' * psi; % Decomposition of initial state psi into the
% M lowest energy levels (cosines and sines)

%a_pdf = abs(a).^2;

%%%%%%%%%%%%%%%%%
%%%% PART B %%%%%
%%%%%%%%%%%%%%%%%

% First, we compare our numerical decomposition coefficients in the
% vector a to the analytical answer a_n = (8/(n*pi))*sin(n*pi/16) if n is
% odd, else a_n = 0.

a_ana = zeros(n_max,1);
if rem(n_max,2) == 0 % that is, if n_max is even
odds = 1:2:n_max-1;
a_ana(odds) = 8*sin(odds*pi/16)./[odds*pi];
else
odds = 1:2:n_max;
a_ana(odds) = 8*sin(odds*pi/16)./[odds*pi];
end

discrep = a_ana - a;
disp('a_n{analytical} - a_n(numerical)')
disp(discrep)

% Second, we plot the initial state as approximated by the lowest n = 2,
% 10, and 20 states

figure
psi2 = phi(:,1:2)*a(1:2);
psi10 = phi(:,1:10)*a(1:10);
psi20 = phi*a;
plot(x,psi.^2,':',x,psi2.^2,'--',x,psi10.^2,'-.',x,psi20.^2,'-');
axis tight
title('Approximation of the Initial State \Psi(x) = 8^{1/2} if |x| < L/16, else 0 by N Lowest Particle-in-a-Box Energy Eigenstates')
xlabel('Position x (in units of box length L)');
ylabel('|\Psi(x)|^2 (probability per unit length)');
legend('\Psi(x,0) analytic','\Psi(x,0) numerical, 2 lowest eigenstates','\Psi(x,0) numerical, 10 lowest eigenstates','\Psi(x,0) numerical, 20 lowest eigenstates');

%cyclelines %% Be careful using cyclelines. Even if you change its poor
%choice of variable name "a" for the figure handles (poor since that's the
%same as our decomposition coefficients) it seems to mess up the later
%waterfall graphs

% Third, fourth, and fifth, we plot the time evolution of psi
% Create our energy and time matrices

m = Me; % we'll say the particle is an electron
Eg = (hbar*pi)^2 / (2*m*L^2); % Ground state energy (n = 1)
varrho = m*L^2/h; % We'll measure time as fractions of the
t = varrho*[0:0.25:1]'; % reconstruction time mL^2 / h

Ecol = [(hbar*pi)^2/(2*m*L^2)]*[1:n_max].^2' ;

E = Ecol*ones(1,length(t));
T = ones(n_max,1)*t';

eta = exp(-i*E.*T/hbar);

a_t = (a*ones(1,length(t))).*eta;

psi2_t = phi(:,1:2)*a_t(1:2,:);
psi2_t_pdf = conj(psi2_t).*psi2_t;
normcheck2 = sum(psi2_t_pdf)*dx;

psi10_t = phi(:,1:10)*a_t(1:10,:);
psi10_t_pdf = conj(psi10_t).*psi10_t;
normcheck10 = sum(psi10_t_pdf)*dx;

psi20_t = phi*a_t;
psi20_t_pdf = conj(psi20_t).*psi20_t;
normcheck20 = sum(psi20_t_pdf)*dx;

%%%%% Time Evolution Figures
%
% NOTE: THESE CAN TAKE A LONG TIME TO PLOT. HENCE, I'VE COMMENTED THEM
% OUT AS DEFAULT. REMOVE THE % SIGNS TO MAKE THE PLOTS.
%
% X_waterfall = x*ones(1,length(t));
% T_waterfall = ones(length(x),1)*t';
%
% figure
% waterfall(X_waterfall',T_waterfall'/varrho,psi2_t_pdf')
% titlestr1 = 'Time Evolution with 2 Lowest Energy Eigenstates';
% titlestr2a = 'Norm Check: \Sigma_x \delta x |Psi(x)|^2 =';
% titlestr2b = num2str(normcheck2(1));
% titlestr2 = [titlestr2a titlestr2b];
% title({titlestr1; titlestr2})
% xlabel('Position x (in units of box length L)')
% ylabel('Time t (in units of the reconstruction time mL^2 / h)')
% zlabel('|\Psi(x,t)|^2 (probability per unit length)')
% axis tight
%
% figure
% waterfall(X_waterfall',T_waterfall'/varrho,psi10_t_pdf')
% titlestr1 = 'Time Evolution with 10 Lowest Energy Eigenstates';
% titlestr2a = 'Norm Check: \Sigma_x \delta x |Psi(x)|^2 =';
% titlestr2b = num2str(normcheck10(1));
% titlestr2 = [titlestr2a titlestr2b];
% title({titlestr1; titlestr2})
% xlabel('Position x (in units of box length L)')
% ylabel('Time t (in units of the reconstruction time mL^2 / h)')
% zlabel('|\Psi(x,t)|^2 (probability per unit length)')
% axis tight
%
% figure
% waterfall(X_waterfall',T_waterfall'/varrho,psi20_t_pdf')
% titlestr1 = 'Time Evolution with 20 Lowest Energy Eigenstates';
% titlestr2a = 'Norm Check: \Sigma_x \delta x |Psi(x)|^2 =';
% titlestr2b = num2str(normcheck20(1));
% titlestr2 = [titlestr2a titlestr2b];
% title({titlestr1; titlestr2})
% xlabel('Position x (in units of box length L)')
% ylabel('Time t (in units of the reconstruction time mL^2 / h)')
% zlabel('|\Psi(x,t)|^2 (probability per unit length)')
% axis tight
%

%%%%%%%%%%%%%%%%%
%%%% PART D %%%%%
%%%%%%%%%%%%%%%%%

% We calculate the mean and standard deviation of the energy in units of
% the ground state energy

mean_energy = Ecol'*[abs(a_t).^2]/Eg;
std_energy = sqrt([(Ecol/Eg).^2]'*[abs(a_t).^2] - mean_energy.^2)


diary off;