%{ 
=========================================================
Solving stationary Schrodinger equation
Method: Finite Difference Approximation (second-order)
Alex G. November 2021

To Do: 
1) plot energy levels together with the potential
=========================================================
%}
clear  % Clears command window
%clc    % Clears command history
clf    % Removes anything in the figure window before simulation. 

tic    % set timer for the program
Date = today('datetime');

fprintf(' Stationary states of Schrodinger equation   \n')
fprintf(' Method: Finite Difference Approximation     \n')
fprintf(' Date:   %s  \n \n', Date)

% Parameters
n    = 2^12;        % number of grid points
if n>2^14           % max 16384 points
  n = 2^14;
end
Key  = 5;           % choosing a potential

nplot= 4;           % nplot first nplot wavefunction
nprint = 12;

fprintf(' Number of grid points  %6i   \n',n)

switch Key
    case 0
        Vp = @V0;            % square well 
        Xmin = -8.00;        % the left-end point
        Xmax =  8.00;        % the right-end point
    case 1
        Vp = @V1;           % infinitely deep square well 
        Xmin = -8.0;        % the left-end point
        Xmax =  8.0;        % the right-end point
    case 2
        Vp = @V2;           % harmonic potential 
        Xmin = -8.0;        % the left-end point
        Xmax =  8.0;        % the right-end point
    case 3
        Vp = @V3;           % a comb 
        Xmin = -10.0;       % the left-end point
        Xmax =  20.0;       % the right-end point
    case 4
        Vp = @V4;           % linear potential 
        Xmin = -10.0;       % the left-end point
        Xmax =  10.0;       % the right-end point    
    case 5
        Vp = @V5;           % Morse potential 
        Xmin =   0.0;       % the left-end point
        Xmax =  10.0;       % the right-end point    
    case 6
        Vp = @V6;           % 1/2 harmonic potential 
        Xmin = -8.0;        % the left-end point
        Xmax =  8.0;        % the right-end point
end

% Format for fprintf
% fmt = [repmat('%12.6f ', 1, nplot), '\n'];

% === call solver
[x,y,Ei,nstates] = QMSch01(Vp,Xmin,Xmax,n);

Ns = min(nstates,nprint);
fprintf(' Energies of bound states  \n')
fprintf(' first %2i out of %4i states \n',Ns,nstates)
fprintf('    i     Energy  \n')

% Print bound states but no more than 10
for k = 1:nstates
    if k > nprint
        break
    end
    fprintf(' %4i %12.6f \n',k, Ei(k,k))
end


% plot(x,y(:,1),'r',x,y(:,2),'b',x,y(:,3),'g')
% title('First three wave functions');
% xlabel('x');
% ylabel('y(x)');

% plot first nplot wave functions
figure (2)
for ip = 1:nplot
    if nstates < ip
        break
    end
    plot (x,y(:,ip),'LineWidth',1.5)
    title('Wave functions');
    xlabel('x');
    ylabel('y(x)');
    hold on
end
hold off


timeE = ceil(toc);
fprintf('\n Elapsed time = %4i sec \n',timeE)

% =========================================================
% THE END of the main part
% =========================================================

function [x,y,Ei,nstates] = QMSch01(Vp,Xmin,Xmax,n)
% =========================================================
% Solver for stationary Schrodinger equation
% Version: November 2021
% =========================================================
%{
IN:
    Vp      - external function (potential)
    Xmin    - left end-point
    Xmax    - right end-point
    n       - number of endpoint
OUT:
    x       - set of grid points
    y       - eigenvectors corresponding to eigenvalies
    Ei      - diagonal matrix of eigenvalues
%}
% 1: preparation
h = (Xmax - Xmin)/(n-1);
a = zeros(n,n);
x = zeros(n,1);
S = zeros(n,1);
Vplot = zeros(n,1);

% 2: Calculate the matrix a(i,j)
% grid points and the diagonal elements
for k = 1:n
    x(k) = Xmin + h*(k-1);
    a(k,k)   = 1.0/h^2 + Vp(x(k));
    Vplot(k)   = Vp(x(k));
end
%  +/- non-diagonal elements
for k=1:n-1
    a(k,k+1) = (-1.0)/(2*h^2);
    a(k+1,k) = a(k,k+1);
end

% 3. Calculating eigenvalues and eigenvectors
% === solver - eigenvalues and eigenvector (Matlab function)
[y,Ei] = eig(a);
% === end solver

% 4. sorting eigenvalues (and eigenvectors) in in ascending order
% attention - Matlab function eig may not always sort
[dwork,ind] = sort(diag(Ei));
EiSort = Ei(ind,ind);
ySort  = y(:,ind);

% testing quality of the function eig.
% e1 = norm(a*y-y*Ei)
% e2 = norm(a*ySort - ySort*EiSort)

Ei = EiSort;
y  = ySort;
% end of sorting

% 5. count bound states
Emax = Vp(Xmax);
nstates = 0;
for k = 1:n
    if Ei(k,k) > Emax
        break
    end
    nstates = nstates+1;
end

% 6. Normalization
% Normalization integral so that |y(x)|^2 dx =1 (method: Simpson)
% note: normalization for the first m functions
for i = 1:nstates
    S(i) = 0.0;
    for k = 2:2:n-1
        S(i) = S(i) + 4.0*y(k,i)*y(k,i);
        S(i) = S(i) + 2.0*y(k+1,i)*y(k+1,i);
    end
    S(i) = S(i) + y(1,i)*y(1,i) - y(n,i)*y(n,i);
    S(i) = sqrt(S(i)*h/3.0);
end
% Normalization
for i=1:nstates
    for k=1:n
        y(k,i) = y(k,i)/S(i);
    end
end

% 7. plot the potential
Vmin = min(Vplot);
Vmax = max(Vplot);
figure (1)
plot(x,Vplot,'k','LineWidth',1.5);
title('Potential V(x)');
xlabel('x');
ylabel('V(x)');
ylim([Vmin-1. Vmax*1.2+1.])
end

% =========================================================
% Square well
% =========================================================
function Vx = V0(x)
a = 2.0;
if abs(x) <= a
    %Vx = -4.0;
    Vx = 0.0;
else
    %Vx = 0.0;
    Vx = 4.0;
end
end

% =========================================================
% Infinitely deep square well
% =========================================================
function Vx = V1(x)
a = 2.0;
if abs(x) <= a
    Vx = 0.0;
else
    Vx = 10000.0;
end

end

function Vx = V2(x)
% =========================================================
% harmonic oscillator
% =========================================================
% a = 6.0;
% if abs(x) <= a
%     Vx = 0.5*x^2;
% else
%     Vx = 0.5*a^2;
% end
Vx = 0.5*x^2;
end


function Vx = V3(x)
% =========================================================
% a comb - set of potential wells
% =========================================================
a  = 1.0;
b  = 0.5;
V0 = -6.0;
if x>=0 && x<=a ...
    || x>=a+b && x<=2*a+b       ...    % n=2
    || x>=2*a+2*b && x<=3*a+2*b ...    % n=3
    || x>=3*a+3*b && x<=4*a+3*b ...    % n=4
    || x>=4*a+4*b && x<=5*a+4*b ...    % n=5
    || x>=5*a+5*b && x<=6*a+5*b ...    % n=6
    || x>=6*a+6*b && x<=7*a+6*b ...    % n=7
    || x>=7*a+7*b && x<=8*a+7*b        % n=8
  Vx = V0;
else
  Vx = 0.0;
end
end


function Vx = V4(x)
% =========================================================
% Wedged linear potential
% =========================================================
% a = 2.0;
% V0= 0.0-2.0*a;
% if abs(x) <= a
%     Vx = V0+2*abs(x);
% else
%     Vx = 0.0;
% end
Vx = 0.5*abs(x);
end


function Vx = V5(x)
% =========================================================
% Morse potential
% =========================================================
a  = 0.5;
V0 = 20.0;
x0 = 1.5;
Vx = V0*(1.0+exp(-2.0*a*(x-x0)) - 2.0*exp(-a*(x-x0)));
end

function Vx = V6(x)
% =========================================================
% 1/2 harmonic oscillator
% =========================================================
a = 6.0;
%if x >= 0.0 && x <= a
if x >=0.0
    Vx = 0.5*x^2;
else
    %Vx = 0.5*a^2;
    Vx = 100.0;
end
end