%****************************************************************
% Program 4: Find several lowest eigenmodes V(x) and
% eigenenergies E of 1D Schrodinger equation
% -1/2*hbar^2/m(d2/dx2)V(x) + U(x)V(x) = EV(x)
% for arbitrary potentials U(x)
% Original at https://arxiv.org/abs/0704.1622
% Modified: R. Herman 10/15/2016
%****************************************************************
clear
close all

% Parameters for solving problem in the interval -L < x < L
% PARAMETERS:
L = 6;                   % Interval Length
N = 1000;                % No of points
x = linspace(-L,L,N)';   % Coordinate vector
dx = x(2) - x(1);        % Coordinate step

hbar = 1; m = 1;          % constants for Hamiltonian

% POTENTIAL, choose one or make your own

%U = 1/2*100*x.^(2);      % quadratic harmonic oscillator potential
%U = 1/2*x.^(4);         % quartic potential

% Finite square well of width a=2w and depth V0
% Pick a/hbar sqrt(2mV0)>npi, for n>1 modes
V0 = 100; 
n = 3;
a = n*pi*hbar/sqrt(2*m*V0);
w = a/2;
H = @(t) zeros(size(t)) + (t >= 0).*(1);
U = V0-V0*(H(x+w)-H(x-w));

% Two finite square wells of width 2w and distance 2a apart
%w = .3; a=3*w;
%H = @(t) zeros(size(t)) + (t >= 0).*(1);
%U = -200*(H(x+w-a) - H(x-w-a) + H(x+w+a) - H(x-w+a));

% Three-point finite-difference representation of Laplacian
% using sparse matrices, where you save memory by only
% storing non-zero matrix elements
e = ones(N,1); 
Lap = spdiags([e -2*e e],[-1 0 1],N,N)/dx^2;

% Total Hamiltonian
H = -1/2*(hbar^2/m)*Lap + spdiags(U,0,N,N);

% Find lowest nmodes eigenvectors and eigenvalues of sparse matrix
nmodes = 3; options.disp = 0;
[V,E] = eigs(H,nmodes,'sa',options);       % find eigs
[E,ind] = sort(diag(E));                   % convert E to vector 
V = V(:,ind);                              % rearrange eigenvectors
if max(V(:,1))< -min(V(:,1))
    V=-V;
end

% Generate plot of lowest energy eigenvectors V(x) and U(x)
set(groot,'defaultAxesColorOrder',[1 0 0;0 1 0;0 0 1],...
      'defaultAxesLineStyleOrder','-|--|:')

Usc = U*max(abs(V(:)))/max(abs(U));        % rescale U for plotting
plot(x,V,x,Usc,'--k');                     % plot V(x) and rescaled U(x)
xlabel('x')
ylabel('U(x),\psi(x)')
title('Potential and Eigenfunctions vs x')

% Add legend showing Energy of plotted V(x)
lgnd_str = [repmat('E = ',nmodes,1),num2str(E)];
legend(lgnd_str) 
shg