%{ 
=========================================================
Solving stationary Schrodinger equation
Method: Shooting method with RK4
Alex G. November 2021
=========================================================
%}

clear  % Clears command window
clc    % Clears command history
clf    % Removes anything in the figure window before simulation. 
tic    % set timer for the program

% Parameters for calculations
Xmin = -8.0;
Xmax =  8.0;
X0   = -1.0;
N    = 1601;
eps  = 1.0e-8;

Emin = -4.0;
Emax = -3.5;

state = 1.0;  %  1.0 for  n = 0,2,4... (even-parity states)
              % -1.0 for  n = 1,3,5,...(odd-parity states)

[Ei] = Qwell02(Xmin,Xmax,X0,N,Emin,Emax,eps,state);
fprintf('\n Ei = %10.6f \n',Ei)

% note for |a|<1, V0=4: Ei = -3.778, -3.120, -2.058, -0.695

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


function [Ei] = Qwell02(Xmin,Xmax,X0,N,Emin,Emax,eps,state)
%{
Solving stationary Schrodinger equation
y'' + 2(E-V(x))y = 0
METHOD: shooting method + Runge-Kutta 4th as initial value problem
CALLS: rk4_2d(x,y,dy,n), ypp(x,y,dy), V(x) (potential)
INPUT: 
  V(X) - a potential (as a function, so far the name is fixed)
  XMIN, XMAX - left and right endpoints for integration
  X0 - "meeting point" (results should not depend on it)
  N  - number of points for the whole interval
  Emin, Emax - energy interval for searching an energy level
  eps - tolerance on matching the log derivative at X0
  state - parity (+1 for odd states and -1 for even states
OUTPUT:
  Ei - eigenvalue (energy)
Additional output (in the function)
Plot 1: matching (left/right) wavefunctions and derivatives
Plot 2: Wavefunction (normalized)
%}
global Ee
MaxIter = 128;            % Max number of iterations

% Working arrays
Delta = zeros(MaxIter,1);
E     = zeros(MaxIter,1);
Vp    = zeros(N,1);

% Phase 1: Prepare arrays x,y xl,yl, xr,yr(left and righ)
h = (Xmax - Xmin)/(N-1);
N0 = ceil((X0-Xmin)/h) + 1;
Nl = N0;
Nr = N - N0 +1;

% reserve arrays
x  = zeros(N,1);
y  = zeros(N,1);

xl  = zeros(Nl,1);
yl  = zeros(Nl,1);
dyl = zeros(Nl,1);

xr  = zeros(Nr,1);
yr  = zeros(Nr,1);
dyr = zeros(Nr,1);

% grid points for x, xl and xr
for k = 1:N
    x(k) = Xmin + h*(k-1);
end
for k=1:N0
    xl(k) = x(k);
end
for k=N:-1:N0
    kr = N-k+1;
    xr(kr) = x(k);
end

% Phase 2: Prepare first three energies (initialization)
E(1) = Emin;
E(2) = Emax;
E(3) = (Emin+Emax)/2;


% Phase 3: the main loop (iterations)

for k = 1:MaxIter
    Ee = E(k);
    b = sqrt(2.0*abs(Ee));
    
%3a Integration from the left to X0
    % initial position and the derivative on the left
    yl(1)  = exp(-b*abs(xl(1)));        
    dyl(1) = b*yl(1);
    [yl,dyl] = rk4_2d(xl,yl,dyl,Nl);

%3b Integration from the right to X0
    % initial position and the derivative on the right
    yr(1) = state*exp(-b*abs(xr(1)));
    dyr(1) = (-1.0)*b*yr(1);
    [yr,dyr] = rk4_2d(xr,yr,dyr,Nr);

%3C Calculating the log derivative difference between the two solutions
    Lelt  = dyl(N0)/yl(N0);
    Right = dyr(Nr)/yr(Nr);
    Delta(k) = (Lelt-Right)/(Lelt + Right);

% print iterations if needed
    if k == 1 % print the header
        fprintf(' Iter   Energy      Match \n')
    end
    fprintf(' %3i   %8.6f  %10.2e \n',k,Ee,Delta(k))
    
    
%3d when k=2 check if there is a solution 
    if k == 2
        if Delta(1)*Delta(2) > 0
            fprintf(' No solution for given Emin and Emax \n')
            Ei = 0.0;
            return
        end
    end
    
    
%3E bisectional search for the root
% Attention: the bisectional method may converge not only to a root 
% of a function f(x)=0 but to a discontinuity, then change the matching
% point X0.
    if k >= 3
        if Delta(1)*Delta(k) < 0.0
            E(2) = E(k);
        else
            E(1) = E(k);
            Delta(1) = Delta(k);
        end
        E(k+1) = (E(1)+E(2))/2.0;
        if abs(Delta(k)) < eps 
%        if abs(E(1)-E(2)) < eps    
            Ei = E(k);
            break
        end
    end   % end bisectional search
               
end % end iterations

if k == MaxIter
    fprintf(' Max iterations - still no solution \n')
    Ei = 0.0;
    return
end
    
figure (1)
plot(xl,yl,'b',xr,yr,'r',xl,dyl,'-m',xr,dyr,'-k','LineWidth',1.5)
title('Matching functions and derivatives')
grid

% Phase 4: Assembling the whole wave function + normalization 

% 4a  One function y(x) as a sum of two solutions (left + right)
Ynorm = yr(Nr)/yl(N0);
for k = 1:N0
    x(k) = xl(k);
    y(k) = yl(k)*Ynorm;
end
for k = 1:Nr
    x(N-k+1) = xr(k);
    y(N-k+1) = yr(k);
end

% 4b  Calculating Integral |y(x)|^2dx for normalization
Sn = 0.0;
for k = 2:2:N-1
    Sn = Sn + 4.0*y(k)*y(k);
    Sn = Sn + 2.0*y(k+1)*y(k+1);
end
Sn = Sn + y(1)*y(1) - y(N)*y(N);
Sn = sqrt(Sn*(h/3.0));
% 4c Normalization
for k = 1:N
    y(k) = y(k)/Sn;
    Vp(k) = V(x(k));
end

figure (2)
plot(x,y,'r','LineWidth',1.5)
%plot(x,y,'r',x,Vp,'b')
str = sprintf('Wave function for Ei = %6.4f',Ei);
title(str);
grid

end %end function


function [y,dy] = rk4_2d(x,y,dy,n)
%{
================================================================
 Solution of a single second-order ODE (Initial-value problem)
 method:  Runge-Kutta 4th-order
 written by: Alex G. (last revision March 2021)
----------------------------------------------------------------
 input ...
  f(x,y,dy)- function from d2y/dx2=f(x,y,dy) (supplied by a user)
  x(i)  - array of x base point (n points)
  y(1)  - initial value for y(1)
 dy(1) - initial value for y'(1)

 output ...
  y(i)  - solutions for y  in n points
 dy(i)  - solutions for y' in n points 
=================================================
%}

for k=2:n
   h   = x(k)-x(k-1);
   k11 = h*dy(k-1);
   k12 = h*ypp(x(k-1),y(k-1),dy(k-1));
   k21 = h*(dy(k-1)+k12/2.0);
   k22 = h*ypp(x(k-1)+h/2.0, y(k-1)+k11/2.0, dy(k-1)+k12/2.0);
   k31 = h*(dy(k-1)+k22/2.0);
   k32 = h*ypp(x(k-1)+h/2.0, y(k-1)+k21/2.0, dy(k-1)+k22/2.0);
   k41 = h*(dy(k-1)+k32);
   k42 = h*ypp(x(k-1)+h,y(k-1)+k31,dy(k-1)+k32);

   y(k) = y(k-1) + (k11 + 2.0*(k21+k31) + k41)/6.0;
  dy(k) = dy(k-1)+ (k12 + 2.0*(k22+k32) + k42)/6.0; 
end
end % function rk4_2d

function [yp2] = ypp(x,y,dy)
% ------------------------------------------
%  Second order ODE as
%  y'' + P(x,y)y' + Q(x,y)y = F(x,y)
%  yp2 = d2y/dx2 = - P(x,y)y' - Q(x,y)y + F(x,y)
% ------------------------------------------
global Ee
P = 0.0;
Q = 2.0*(Ee - V(x));
F = 0.0;
yp2 = (-1.0)*P*dy +(-1.0)*Q*y + F;
end

function Vx = V(x)
a = 2.0;
if abs(x) <= a
    Vx = -4.0;
else
    Vx = 0.0;
end

end