% (0) modified on 2015-8-10, running on Matlab R2014b
% (1) This code can be used to simulate the AC Stark shift and
% post-ionization excitation on various ionics states, more in readme.
% (2) Two pulses are used: Low frequency and high intensity IR, mimicing
% the real pulse in the lab. High frequency and low intensity vuv (labeled as X),
% coupling the neutral ground to the desired ionic state.
% (3) An example of 5 states in the diabatic basis is given
% 

function [w, mu, t, a, ProbAmp, Pop] = IwPE_template()

%Define parameters
tic
global hbar hbar_ev auDipole w mu EX EIR omegaX omegaIR tIR tX Power Basis NStates

Power = 12; %power factor used to select unit, in this case, "Power=12" gives unit of ps
hbar = 1.054571726e-34*10^Power; % [J ps]
hbar_ev = 6.58211928e-16*10^Power; % [eV ps]
Epsilon = 8.8541878176e-12; % [F/m]
Hartree = 4.35974417e-18/1.6021765e-19; % = auEnergy, [eV]
auDipole = 8.47835326e-30; % [C m]
auField = 5.1422065e11; % [V/m]
NStates = 5;   %number of state
w = []; % Potential energy/frequency matrix
mu = []; % transition dipole moment (TDM) matrix

% pulse parameters 
EX = 0.00217*auField; % [V/m], Peak int = EX^2*Epsilon*3e8/2 [W/m^2]
% omegaX = 10.37/hbar_ev; omegaX can be set here arbitrarily or later based on the potential energies of states
EIR = 0.0237*auField; % 0.0237au~20 TW/cm^2, 10e16 W/m^2~0.0169 au, 100e16W/m^2~0.0534 au
omegaIR = 0.05843*Hartree/hbar_ev; % [radian/ps], IR wavelength = 3e8/(omegaIR/2/pi*1e12) [m]

% 'tIR' and 'tX' are characteristic times for IR and VUV pulses,depending on the pulse shape. E.g.,
% Gaussian: exp(-1*(t.^2)/2/tX^2) --> FWHM_field=2*sqrt(2*log(2))*tX=2.35482*tX, FWHM_int=2sqrt(ln2)tX=1.6651*tX [ps];
% cos^2: (cos(pi*t/tX))^2.*(t>-tX/2 & t<tX/2) -->FWHM_field=0.5*tX,FWHM_int=2*acos(0.5^0.25)/pi*tX=0.364*tX [ps]
tIR = 30e-3/2*pi/acos((.5)^0.25); % [ps], cutoff time for cos^2 waveform
% tX = 20e-3/2*pi/acos((.5)^0.25); %20e-3/2*pi/acos((.5)^0.25); %ps in cos^2(pi*t/tX) -->FWHM_int=20 fs
tX = tIR/7^.5;

% choose which basis to used, partially depending on the energy and TDM data
% given. One could transform from diabatic to adiabatic. The opposite
% direction doesn't seem useful.
Basis = 'd';
% Basis = 'a';
%%
% Population transfer calculation
if Basis == 'd' % Energies and SOC in diabatic basis, directly from Tamas' calculation
    display 'diabatic'
%     Energies
    w_di = zeros(NStates); SOC = zeros(NStates);
    w_di(1,1) = -0.3358; % [au], neutral ground
    w_di(2,2) = 0.0079; % D0(+)
    w_di(3,3) = 0.0079; % D0(-)
    w_di(4,4) = 0.0141; % D1(+)
    w_di(5,5) = 0.0141; % D0(-)
    w_di = w_di*Hartree/hbar_ev; % [rad/ps]
    SOC(2,5) = -0.0085 + 1i*0.0062; % spin-orbit coupling, [au]
    SOC(5,2) = -0.0085 - 1i*0.0062;
    SOC(3,4) = 0.0085 + 1i*0.0062;
    SOC(4,3) = 0.0085 - 1i*0.0062;
    SOC = SOC*Hartree/hbar_ev; % [rad/ps]
    w_SOC = w_di + SOC;
%     Diagnolization of (bare+SOC) Hamiltonian. Apply the similarity
%     transformation to get the populatoni in the adiabatic basis (the
%     eigenbasis defined by bare+SOC Hamiltonian).
    [v, w_ad]=eig(w_SOC); % diagnolization, ie., w_ad = v\w_SOC*v
%     u = inv(v); % use either u*a or v\a to transform the final state vector 'a' to adiabatic basis.
    w = w_SOC; % used in solving TDSE
    
%     Transition dipole moments, TDM
    TDM_error = 1;
    mu_di = zeros(NStates);
    mu_di(2,4) = -1.875e-8; % [au]
    mu_di(4,2) = -1.875e-8;
    mu_di(3,5) = -1.875e-8;
    mu_di(5,3) = -1.875e-8;
%     diagonal entries are generally non-zero but not used in the
%     calculation. The TDMs between the neutral and the ionic are assumed
%     to be 0.1 au since they are difficult to calculate and not to our
%     interest.
    mu_di(1,2)=1/sqrt(2);    mu_di(2,1)=1/sqrt(2);
    mu_di(1,3)=1/sqrt(2);    mu_di(3,1)=1/sqrt(2);
    mu_di = mu_di*auDipole;
    mu = mu_di*TDM_error; % used in solving TDSE
%     mu_ad = v\mu_di*v;
    
    omegaX = w(2,2) - w(1,1); % VUV set to be on resonance
    [t,a] = SolveTDSE_IwPE(); %solve TDSE
    
    for j = 1:NStates
        a(:,j) = a(:,j).*exp(-1i*w(j,j)*t);
    end
    clear ProbAmp Pop;
%     IMPORTANT: for transpose, use .' instead of '
    ProbAmp = v\a.'; % Population amplitude in adiabatic basis
    Pop = abs(ProbAmp.').^2;  %population
    Pop(:,end+1) = t;  
end
%%
if Basis == 'a' % adiabatic basis
    display 'adiabatic'
    
    w_ad = zeros(NStates);
    w_ad(1,1) = -0.3358; % [au], neutral ground
    w_ad(2,2) = 3.186e-5;
    w_ad(3,3) = 3.186e-5;
    w_ad(4,4) = 0.022;
    w_ad(5,5) = 0.022;
    w_ad = w_ad*Hartree/hbar_ev; % [rad/ps]
    w = w_ad;
    
    TDM_error = 1;
    mu_ad = zeros(NStates);
%     data for adiabitic TDM goes here ...
    mu = mu_ad*auDipole*TDM_error;
    
    omegaX = w(2,2) - w(1,1);
    [t,a] = SolveTDSE_IwPE(); %solve TDSE
    Pop = abs(a).^2;  %population
    Pop(:,end+1) = t;
    ProbAmp(r,:) = a(end,:);
%     PopE(p,r,:) = abs(ProbAmp(r,:)).^2;      
end
toc
end

function [t,a] = SolveTDSE_IwPE()
global  tIR tX NStates mu
options = odeset('RelTol',1e-6,'AbsTol',1e-7); % ode solving accuracy
a0 = zeros(NStates,1)'; % ket vector
a0(1) = 1; % initial condition, sum of amplitude square should be 1
% a0(4) = 1/sqrt(2);
% a0(5) = 1/sqrt(2);
[t,a] = ode23(@TDSE_IwPE,[-2*max(tX,tIR) 2*max(tX,tIR)],a0,options); % set solver, time interval, initial condition.
end

function da = TDSE_IwPE(t, a)
global hbar w mu EX EIR omegaX omegaIR tIR tX Basis NStates
% t in ps

MU = zeros(NStates);
for j=1:NStates
    %     use RWA for VUV coupling the neutral ground and the ionic ground. 
        MU(1,j) = 1/2/hbar*mu(1,j)*(EX*cos(pi*t/tX).^2.*(t>-tX/2 & t<tX/2)*exp(+1i*omegaX*t)+2*EIR*cos(pi*t/tIR).^2.*(t>-tIR/2 & t<tIR/2)*cos(omegaIR*t))*exp(-1i*(w(j,j)-w(1,1))*t);
        MU(j,1) = 1/2/hbar*mu(j,1)*(EX*cos(pi*t/tX).^2.*(t>-tX/2 & t<tX/2)*exp(-1i*omegaX*t)+2*EIR*cos(pi*t/tIR).^2.*(t>-tIR/2 & t<tIR/2)*cos(omegaIR*t))*exp(-1i*(w(1,1)-w(j,j))*t);
    %     no RWA, no noticeable differenct
%     MU(1,j) = 1/2/hbar*mu(1,j)*(EX*cos(pi*t/tX).^2.*(t>-tX/2 & t<tX/2)*2*cos(omegaX*t)+EIR*cos(pi*t/tIR).^2.*(t>-tIR/2 & t<tIR/2)*2*cos(omegaIR*t))*exp(-1i*(w(j,j)-w(1,1))*t);
%     MU(j,1) = 1/2/hbar*mu(j,1)*(EX*cos(pi*t/tX).^2.*(t>-tX/2 & t<tX/2)*2*cos(omegaX*t)+EIR*cos(pi*t/tIR).^2.*(t>-tIR/2 & t<tIR/2)*2*cos(omegaIR*t))*exp(-1i*(w(1,1)-w(j,j))*t);
end
if Basis == 'd';    % diabatic basis, with SOC matrix
    for k=2:NStates
        for j =2:NStates
            if j~=k
                MU(k,j) = 1/hbar*mu(k,j)*EIR*cos(pi*t/tIR).^2.*(t>-tIR/2 & t<tIR/2)*cos(omegaIR*t)*exp(-1i*(w(j,j)-w(k,k))*t) - w(k,j)*exp(-1i*(w(j,j)-w(k,k))*t);
            end
        end
    end
elseif Basis == 'a'     % adiabatic basis, no SOC H, approx, matrix form
    for k=2:NStates
        for j =2:NStates
            if j~=k
                MU(k,j) = 1/hbar*mu(k,j)*EIR*cos(pi*t/tIR).^2.*(t>-tIR/2 & t<tIR/2)*cos(omegaIR*t)*exp(-1i*(w(j,j)-w(k,k))*t);
            end
        end
    end
end
da = 1i*MU*a;

end
