% *******************************************
% Weinacht Group Strong Field Ionization Code
% *******************************************
% *******************************************
% This code solves for the time-dependent state
% populations & photoelectron spectrum for a
% femtosecond laser pulse driving a generic system
% with a single ground state, single excited bound
% state, and a single ionic state:
% 
%          /////// Continuum
%             ^
%             |   (n12 photons)
%             |
%           ----- |1>
%             ^
%             |   (n01 photons)
%             |
%           ----- |0>
% 
% *******************************************
% This code intended as a bare-bones working TEMPLATE.
% Comments throughout indicate where and how to extend
% the code to model more complicated systems.
% *******************************************
function output = StrongFieldIonization
% SECTION I: System and Laser Parameters (SI units, all times in ps)

% A: Constants
HBAR = 1.054572E-34; %Reduced Planck's constant
ALPHA = 1/(2*HBAR*2)*1e-12; %Coupling parameter for continuum ladder
EMAX = 3.2e-19; %Maximum photoelectron energy [J] (i.e., cutoff)
M = 400; %Number of states in continuum ladder
NUMBER_OF_BOUND_STATES = 2; %Must match number of bound states 

% B: Laser
wField = 2*pi*384; %Central frequency of laser pulse
FWHM = 0.04; %FWHM of laser pulse
Io = 5.0E016; %Intensity of laser pulse [W/m^2] 
I_REF = 1.278E017; %Reference intensity to which peak Rabi frequencies and Stark shifts (defined below) correspond
% (Add e.g. pulse shape parameters here)

% C: Levels
w0 = 0; %Peak AC Stark shift of level |0>
w1 = 2*pi*176; %Peak AC Stark shift of level |1>
wC = 2*pi*176; %Peak AC Stark shift of Continuum
% The 176 number comes from the ponderomotive shift at the reference
% intensity.  The ponderomotive shift is ~14.4 THz/(TW/cm^2) and the ref
% intensity is given above
% (ADD HERE Stark shifts for any additional levels)

% D: Transitions
%   Notation is (QUANTITY)(Lower Level)(Upper Level)
%   w = transition frequency ; n = photon number;
%   rabi = peak amplitude of transition matrix element
%   e.g. rabi01 is rabi frequency for |0>->|1> transition
n01 = 5;
w01 = n01*wField - 2*pi*80; % used 100 for paper calcs
rabi01 = 2*pi*10.22; %  used 1.22 for paper calcs

n1C = 2;
w1C = n1C*wField - 254*2*pi + EMAX/HBAR/2*1E-12; %(See paper for info on the EMAX term)
rabi1C = 2*pi*71;
% (ADD HERE parameters for additional transitions) 

% *******************************************

% % II: Laser Pulse Shaping

% A: Define time axis from -300fs to +300fs
step=0.001;
lim = 1; 
times=-lim:step:lim;

% B: Define an unshaped Gaussian pulse with specified FWHM
tau = 2*FWHM/2.35;
e_t=exp(-times.^2/tau^2);

% C: Fourier transform the pulse
e_t_shift=fftshift(e_t);
e_f=fft(e_t_shift);

% D: Apply frequency-domain SHAPING, if desired 
f=-1/(2*step):1/(2*lim):1/(2*step); %Define frequency axis, centered about f = 0
phase=0*fftshift(pi/2*sign(f-25)); %Define a spectral phase profile; here, a pi phase jump at the central frequency
e_f_shaped=e_f.*exp(-1i*phase); %Apply spectral phase profile

% E: Transform shaped pulse back to time domain
e_t_shaped=ifft(e_f_shaped);
timePulse = fftshift(e_t_shaped); %Complex Electric Field vs. Time


% *******************************************

% III: Hamiltonian (See Paper for Discussion)

% This function represents the RHS of the TDSE written as:
%    Dt(s) = -i/hbar . H . s
% but is loosely named "Hamiltonian" in this code

function sp=Hamiltonian(t,s)

N = M + NUMBER_OF_BOUND_STATES; %Number of rows in Hamiltonian

% Preliminary: Interpolation of the Pulse E Field
ii=find(times>t,1);
dt=times(ii)-times(ii-1);
g=(abs(timePulse(ii))^2-abs(timePulse(ii-1))^2)/dt*(t-times(ii-1))+abs(timePulse(ii-1))^2;%Interpolating the |E|^2 values
phi=(angle(timePulse(ii))-angle(timePulse(ii-1)))/dt*(t-times(ii-1))+angle(timePulse(ii-1));%Inperpolating the ARG(E) values

% A: First, set continuum ladder couplings in EVERY row

sp=-1i*wC*g*s*Io/I_REF;
sp(1:(N-1))=sp(1:(N-1))-1i*EMAX*ALPHA*s(2:N);
sp(2:N)=sp(2:N)-1i*EMAX*ALPHA*s(1:(N-1));

% B: Overwrite bound state rows, leaving ladder couplings everywhere else

sp(1:3) = -1i*[s(1)*w0*g*Io/I_REF + s(2)*rabi01*(Io*g/I_REF)^(n01/2)*exp(-1i*(w01 - n01*wField)*t)*exp(-n01*1i*phi),...
s(2)*w1*g*Io/I_REF + s(1)*rabi01*(Io*g/I_REF)^(n01/2)*exp(1i*(w01 - n01*wField)*t)*exp(n01*1i*phi) + rabi1C*(Io*g/I_REF)^(n1C/2)*exp(-1i*(w1C - n1C*wField)*t)*s(3)*exp(-n1C*1i*phi),...
s(3)*wC*g*Io/I_REF + rabi1C*(Io*g/I_REF)^(n1C/2)*exp(1i*(w1C - n1C*wField)*t)*s(2)*exp(n1C*1i*phi) + s(4)*EMAX*ALPHA,...
];

% ADD HERE any changes to couplings, new couplings, new levels, etc.


end
% *******************************************

% IV: Numerical Integration

% A: Set Desired Precision for Integration
PREC_ABS = 5E-6; %Absolute Precision
PREC_REL = 5E-6; %Relative Precision

% B: Set Initial State Vector and Call Integrator

systemSize = NUMBER_OF_BOUND_STATES + M;
s0=zeros([systemSize 1]); %Create initial state vector
s0(1)=1; %All population initially in ground state

T_LIMIT = 0.1; %Limits of integration are +/-T_LIMIT
tic
options = odeset('RelTol',PREC_ABS,'AbsTol',PREC_REL); %Options for integrator
[t,s] = ode23(@Hamiltonian,[-T_LIMIT T_LIMIT],s0,options); %Call integrator

% *******************************************

% V: Photoelectron Spectrum

% A: Select the Continuum Coefficients at Some Time

DESIRED_TIME = 0.05; %(Time at which to compute photoelectron spectrum)
s_ion=s(:,NUMBER_OF_BOUND_STATES+1:(M+NUMBER_OF_BOUND_STATES)); %Extract the continuum state coefficients from the state vector
CLOSEST_TIME = find(t > DESIRED_TIME,1); %(Picking the timestep closest to desired time) 
ionPlot = s_ion(CLOSEST_TIME,1:M); %The array of continuum coefficients at CLOSEST_TIME

% ADD HERE a similar selection of coefficients for any additional continua

% B: Generate Photoelectron Spectrum Using Legendre Polynomials
eMaxReduced = EMAX/(1.602e-19); %Write EMAX in eV
resolution = 0.01; % [eV]
energies=0:resolution:eMaxReduced; % Create energy axis for the calculation

% For an explanation of the lines below, see the manuscript and
% references therein:

TEMP = size(ionPlot); 
MAX_L = TEMP(2);
runningSum = 0*double(energies);

for l=1:1:MAX_L
    legendAll = legendre(l-1,2*energies/eMaxReduced - 1);
    polynomial = legendAll(1,:);
    runningSum(:) = ionPlot(l)*sqrt((2*l+1)/eMaxReduced)*polynomial' + runningSum(:);
    
end

PHOTOELECTRON_SPECTRUM = abs(runningSum).^2; %The photoelectron spectrum

% *******************************************

% VI: Display Results

% A: Plot State Populations vs. Time
% NOTE: t is array of timesteps from numerical integration
%       times is array of timesteps from pulse shaping

% Preparing axis limits:
TI1 = find(times>-T_LIMIT,1);
TI2 = find(times>T_LIMIT-step,1);

subplot(2,2,1)
plot(t,abs(s).^2);
hold on;
plot(t,sum(abs(s_ion').^2),'black'); % Plotting Total Ionized Fraction in Black
plot(times(TI1:TI2),abs(timePulse(TI1:TI2)).^2,'red'); %Plotting Pulse Intensity Profile in Red
legend('|0>','|1>');
titlestr = 'State Populations vs. Time';
title(titlestr);
xlabel('Time from Pulse Peak [ps]');
ylabel('Population');
hold off;

subplot(2,2,2)
plot(t,abs(s(:,2)).^2,'green');
legend('|1>');
titlestr = 'Population of |1> vs. Time';
title(titlestr);
xlabel('Time from Pulse Peak [ps]');
ylabel('Population');

subplot(2,2,3)
plot(t,abs(s_ion).^2);
titlestr = 'Population of Ladder Coefficients vs. Time';
title(titlestr);
xlabel('Time from Pulse Peak [ps]');
ylabel('Population');

% B: Plot Photoelectron Spectrum
subplot(2,2,4)
plot(energies,PHOTOELECTRON_SPECTRUM);
titlestr = 'Photoelectron Spectrum';
title(titlestr);
xlabel('Energy [eV]');
ylabel('Photoelectron Yield (Arb.)');

toc
end

              






