%------------------------------------------------------------------------%

%---Solution of Time Independent Schrodinger Equation using Finite

Difference---------%
%------------------for an artibrary potential-------------------------------%
%------------------------------------------------------------------------%

clc
close all
clear all

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%% Constants %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

h=6.62606896E-34; %% Planck constant [J.s]

hbar=h/(2*pi);
e=1.602176487E-19; %% electron charge [C]
kB = 1.3806488E-23; %% Boltzmann's constant [J/K]
T=300; %% Temperature [K]
m0=9.10938188E-31; %% electron mass [kg]
me=0.1; %% Effective mass of electron in the well
Ef=0.1; %% Fermi level
n=5; %% Number of desired eigen energy levels
ScF= 0.1;% scaling factor to plot the wave function [Without Dimension]
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%% Potential Energy Function %%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Grid %%%
Nz=101;
dz=1e-10; % m
z=0:dz:(Nz-1)*dz;

%%% potential definition %%

U=zeros(1,Nz);
U0=2;
Zmax=(Nz-1)*dz;
slop = (U0/Zmax);
U=slop*z

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%% Building of the operators %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

DZ2 =(-2)*diag(ones(1,Nz)) + (1)*diag(ones(1,Nz-1),-1) + (1)*diag(ones(1,Nz-1),1);

DZ2=DZ2/dz^2;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%% Total Hamiltonian %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

H = -hbar^2/(2*m0*me) * DZ2 + diag(U*e) ;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%% Calculation of Eigen Energies %%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

H = sparse(H);
[psi,Energy] = eigs(H,n,'SM');
E = diag(Energy)/e ;
E=real(E);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%% Normalization of the Wavefunction %%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

for i=1:n
psi(:,i)=psi(:,i)/sqrt(trapz(z',abs(psi(:,i)).^2)); % normalisation at 1
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%% Scaling and shifting the wavefunctions %%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

for i=1:length(E)
PSI(:,i)=abs(psi(:,i)).^2/max(abs(psi(:,i)).^2)*ScF + E(i); % normalisation
for the plotting
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%% Charge density calculation %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Here, I define the energy grid for the charge density calculation
dE=1e-4; % Resolution of the energy grid
En = E(1):dE:max(U); % Range of the energy grid

ro=[];
for i=1:length(E)
ro( En>E(i),i) = (m0*me)/(pi*hbar^2); % Write the prefactor here
ro( En<E(i),i) = 0;
end

roEf=[];
Fermi= 1./(1+exp((En-Ef)/(kB*T/e))) ;
for i=1:length(E)
roEf(:,i) = ro(:,i).*Fermi'; % Charge density with energy for each sub-band
NN(i)= sum(roEf(:,i)); % Charge density for each sub-band by summing the
charge density over energy
end

ntot = repmat(NN,[length(z) 1]).*abs(psi).^2 ; % Charge density with distance for

each sub-band

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%% For plotting the charge density with energy,distance and sub-
band %%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%

for ii=1:length(E)
psic_M = repmat(psi(:,ii),[1,length(En)])';
ROEf(:,:,ii) = repmat(roEf(:,ii),[1 length(z)]) .* abs(psic_M.^2);
end
figure(1)
pcolor(z*1e9,En,sum(ROEf,3)*1e-6 )
shading flat
hcb=colorbar;
hold on
plot(z*1e9,U, 'b','linewidth',1)
hold on
for i=1:length(E)
plot(z*1e9,PSI(:,i),'color','r','linewidth',1)
end

xlabel('z (nm)')
ylabel('Energy (eV)')