%--------------------------------------------------------------------------
%
%
%    Program to calculte the energy levels for a particle-in-a-box
%    using a matrix based method.
%
%    Marcel Drabbels, 28-9-2018
%
%--------------------------------------------------------------------------

%------------ DEFINE CONSTANTS --------------------------------------------
h_c=6.626070040e-34;                        % Planck's constant
h_bar=h_c/2/pi;                             % Planck's constant h-bar 
me=9.10938356e-31;                          % mass electron
mp=1.6726219e-27;                           % mass proton
mn=1.674929e-27;                            % mass neutron
amu=1.660540e10-27;                         % atomic mass unit
eV=1.6021766208e-19;                        % electron volt in SI units
Lse=1e-10;                                  % Angstrom in SI units
bound=0;                                    % number of bound states


%------------ SET PARAMETERS FOR CALCULATIONS -----------------------------
N=501;                                      % Number of data points. NOTE: number of intervals equals N-1
xmin=-0.5;                                  % Minumum x value for the calculation
xmax=+1.5;                                  % Maximum x value for the calculation
m=me;                                       % mass for calculation
a=1.0;                                      % size of box (in Angstroms)
V0=1.0*10^12;                               % potential outside box (in eV), V0=10^12 for infinite


%------------ SETUP MATRICES ----------------------------------------------
SD_matrix = zeros(N);                       % Initialize matrix for second derivative
K_matrix = zeros(N);                        % Initialize matrix for kinetic energy
V_matrix = zeros(N);                        % Initialize matrix for potential
U=0;                                        % Potential array

dx=(xmax-xmin)/(N-1);                       % Calculate x-interval
xx = linspace(xmin,xmax,N);                 % Create x-vector for plotting potential and wavefunctions
xl(1)=0;                                    % Create x-vector for plotting energy level
xl(2)=a;                                    % Create x-vector for plotting energy level

Cke=-h_bar^2/2/m/(eV*Lse^2);                % Constant for kinetic energy operator

for i=1:N                                   % Loop over position
    SD_matrix(i,i)=-2;                      % Set diagonal elements SD matrix   
    SD_matrix(i+1,i)=1;                     % Set off-diagonal elements SD matrix
    SD_matrix(i,i+1)=1;                     % Set off-diagonal elements SD matrix
    x=xmin+(i-1)*dx;                        % Calculate x position (Angstrom)
    if (x <= 0) || (x >= a)
        U(i)=V0;                            % Set potential matrix outside the box to V0
    else
        U(i)=0;                             % Set potential matrix inside the box to 0
    end
    V_matrix(i,i)=U(i);                     % Set Set diagonal elements V matrix   
    E(i)=h_c^2/(8*m*a^2)*i^2/eV/Lse^2;      % Calculate exact values (in eV) of lowest 10 energies for a particle-in-a-box for comparison later on
end

K_matrix=Cke/dx^2*SD_matrix;                % Calculate kinetic energy matrix
H_matrix=K_matrix(1:N,1:N)+V_matrix;        % Calculate Hamiltonian matrix


%------------ DIAGONALIZE MATRIX ------------------------------------------
tic                                         % Get timing informataion
[Evec,Eval]=eig(H_matrix);                  % Find eigenvalues and normalized eigenfunctions of Hamiltonian using the Matlab routine eig
toc                                         % Get timing informntaion


%------------ WRITE RESULTS TO SCREEN -------------------------------------
disp('Quantum State       Eigenvalue              Exact Energy            Energy difference       Rel. energy difference  ');
for n = 1 : 10
    fprintf('    %0.0f   \t     ',n);                           % Quantum state
    fprintf('   %#5.8f \t     ',Eval(n,n));                     % Calculated eigenvalue
    fprintf('   %#5.8f \t     ',E(n));                          % Exact eigenvale
    fprintf('   %#5.8f \t     ',Eval(n,n)-E(n));                % Calculate absolute deviation from exact result
    fprintf('   %#5.8f \t  \n',(Eval(n,n)-E(n))/E(n));          % Calculate relative deviation from exact result 
end


%------------ PLOT RESULTS ------------------------------------------------
figure(1);
clf                                                             % Clear figure
set(gcf,'Name','Potential & Energy levels');                    % Set name figure
xlabel('x (Angstrom)');                                         % Set x label
ylabel('Energy (eV)');                                          % Set y label
grid on                                                         % display grid
hold on                                                         % Set multiple graphs
if (V0 > 1e8)                                                   % Scale axis
    axis([xmin xmax  min(U)-(E(10)-min(U))*0.1   min(U)+(E(10)-min(U))*1.1]); 
else
    axis([xmin xmax  min(U)-(max(U)-min(U))*0.1   min(U)+(max(U)-min(U))*1.1]);
end
plot(xx,U,'LineWidth',3);                                       % Plot potential

for n=1:N
    yl(1)= Eval(n,n);                                           % set y vector energy levels
    yl(2)= yl(1);                                               % set y vector energy levels
    if (Eval(n,n) <= V0)                                        % Plot bound states
        plot(xl, yl, 'r','LineWidth',1);
    end
end


figure(2);
clf
set(gcf,'Name','Eigenfunctions & Prob. densities');
for cn = 1:3
    subplot(2,3,cn);
    y1 = 0.5*(1+0.75*Evec(:,cn) ./ (max(Evec(:,cn))));
    y2 =  (U-min(U)) ./ (max(U) - min(U));
    hold on
    axis off
    title_m = ['\psi   n = ', num2str(cn)] ;
    title(title_m,'Fontsize',12);
    plot(xx,y2,'lineWidth',3)
    plot(xx,y1,'r','lineWidth',2) 
    
    subplot(2,3,3+cn);
    y1 = 0.5*(1+0.75*Evec(:,cn).^2 ./ (max(Evec(:,cn).^2)));
    y2 =  (U-min(U)) ./ (max(U) - min(U));
    hold on
    axis off
    title_m = ['\psi^2   n = ', num2str(cn)] ;
    title(title_m,'Fontsize',12);
    plot(xx,y2,'lineWidth',3)
    plot(xx,y1,'g','lineWidth',2)

end

