%function p_final=i2_chirp_movie( beta_w)   
% Wavepacket focussing using chirp in diatomic iodine (I2).  Computes the
% first order wavefunction as a function of time for single Gaussian pulses of
% width FWHM with chirp beta_w in wavenumbers/femtosecond.  The carrier
% frequency w0 is given as a function of time
% w0(t)= w0_1 + beta_w*(t-t0_1)
% where w0_1 and t0_1 are the center frequency and time of the maximum
% intensity of the pulse respectively.
% If plot_wave>0, the first order wavefunction (absolute value in green, real
% part in red) is plotted during the wavepacket propagation.  At the end of 
% the propagation, the electric field intensity, the excited state population, 
% and the expectation value of the bond length in the excited state are
% plotted as a function of time in the second window.  The function returns
% the final population in the excited state. 
%
% The first order function Wigner representation is calculated and displayed.
% A movie of the evolution is displayed.
%
% This program also calculates and displays the location of the center of
% the first order w.f. This can then be used to design an optimal excitation
% pulse (by finding omega(x(t)) - the frequency difference between the energy
% surfaces along the trajectory (see exercise 7).

% define various conversion constants

clear;

AU_TO_EV=27.2113961;
NM_TO_AU=45.56335;
AU_TO_CM=2.194746e5;
AU_TO_ANG=0.529177249;
AU_TO_FS=0.0241889;
AMU_TO_AU=1.8228885e3;
conv_r= AU_TO_ANG;
conv_en=AU_TO_CM;
conv_t= AU_TO_FS;

% Turn on and off plotting of the wavepacket wavefunction
plot_wave=-1;

% set up coordinate grid
rmin= 4.3; % Note: in Bohr
rmax=  9.0; 
%ngrid=128; % Should be power of 2 for good performance of FFT
ngrid=256;%---------------------------------------
dr= (rmax-rmin)/(ngrid-1);
r= rmin:dr:rmax;

% mass of I2
mass= (126.904473/2)*AMU_TO_AU;

% X (ground) state of I2---not used here for dynamics but useful for
% visualization 
De_X= 12440/AU_TO_CM;
be_X= 1.875*AU_TO_ANG;
re_X= 2.656/AU_TO_ANG;

V1= De_X*(1-exp(-be_X*(r-re_X))).^2;

% B (excited) state of I2
De_B=  5169/AU_TO_CM;
be_B= 1.696*AU_TO_ANG;
re_B= 3.025/AU_TO_ANG;
Te_B= 15771/AU_TO_CM;

V2= De_B*(1-exp(-be_B*(r-re_B))).^2 + Te_B;

% initial ground state wavefunction (ground state of harmonic oscillator---an
% approximation here, but a good one.)
w_vib= (210)/AU_TO_CM;
beta= -sqrt(mass*w_vib);
psi_init= exp(-0.5*beta^2*((r-re_X).^2));
E_init= w_vib/2;
norm= 1/sqrt(dot(psi_init, psi_init));
psi_init= norm*psi_init;

% pulse parameters
FWHM=  50.0/AU_TO_FS; % This is FWHM of pulse intensity, not amplitude
w0_1= 17000/AU_TO_CM; % Pulse 1 frequency
beta_w= beta/AU_TO_CM*AU_TO_FS;  %--------------------------------------------- 
t0_1=  2*FWHM; % center of laser pulse 1
% Pulse intensities in atomic units
e0_1= 0.00386;     
% Transition dipole in atomic units
mu12= 0.181;    

% integration parameters
dt      =     0.1 /AU_TO_FS;
t_sample=    10   /AU_TO_FS;
%t_end   =   t0_1 + 600/AU_TO_FS;
t_end   =   t0_1 + 308/AU_TO_FS;
% t_end=3*t_end;%-- the orig. t_end -------------------------------------
t_end=9*t_end;%------ For 3 piriods ----------------------------------------
update_freq= 1; % update_freq*t_sample is the interval at which plots updated
n_steps_sample= round(t_sample/dt);
n_sample= round(t_end/t_sample);
t_end= n_sample*t_sample;

% pulse function
t= 0:dt:t_end;
alpha= 2*log(2)/FWHM^2;
w0= w0_1 + beta_w*(t-t0_1);
E= e0_1*exp(-alpha*(t-t0_1).^2).*exp(-i*w0.*(t-t0_1));

cpu_start= cputime;
clock_start= clock;

% set up k grid
dk= 2*pi/(ngrid*dr);
k=[[0:ngrid/2-1],[-ngrid/2:-1]]*dk;

% set up kinetic energy and k-space propagator
kin= (0.5*(k.^2)/mass).';

% set up propagators
u_kin= exp(-i*dt*kin);
u_kin_half= exp(-i*0.5*dt*kin);
u_pot1= exp(-i*dt*(V1.'));
u_pot2= exp(-i*dt*(V2.'));

psi_0=    psi_init.';          % unperturbed wavefunction
psi_1=    zeros(ngrid, 1);     % First order wavefunction%--------------------------------

% Plot potential curves and initial wavefunction
if (plot_wave)
   figure(1);
   subplot(1,1,1);
   clf;
   pot_plot=plot( r*conv_r, [V1*conv_en; V2*conv_en]);
   axis manual;
   axis([conv_r*min(r) conv_r*max(r) 0 conv_en*V2(ngrid)]);
   hold on;
   wave_scale= 0.2*De_X*conv_en/(max(psi_init)-min(psi_init));
   psi_plot=plot( r*conv_r, wave_scale*psi_init+ E_init*conv_en, ...
                  r*conv_r, 0*psi_init+ (E_init+w0_1)*conv_en,...
                  r*conv_r, 0*psi_init+ (E_init+w0_1)*conv_en);
        
   %print -djpeg -r72 'i2_chrp_mv1.jpeg';  %-------------------------------------------           
          end

tot_steps=0;
t=0;
max_steps= round(t_end/dt);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%initialize movie making!
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%W=Wig(psi_1);
W=wig2((conj(psi_1))',r);
figure(2);
surf(W);
shading interp;
view(0,80);
axis ([0 250 0 300 -5 5]) ;     %axis([0 128 0 128 -0.05 0.05]);
caxis(.5*[(1.0e-008)*-.5  (1.0e-006)*0.4]);%----------------------------------------------------
M=moviein(n_sample); %---------------------------------------
%%%%%%%%%%%%%%%%%%%%

n_sample=0;

while (tot_steps<max_steps)
   % propagate wavefunctions 1/2 step forward
   U= u_kin_half;
   psi_1=    ifft(U.*fft(psi_1));

   % Propagate until next sample time 
   for step=1:n_steps_sample
      % potential step
      psi_0=    psi_0*exp(-i*dt*E_init);
      psi_1=    u_pot2.*psi_1;
      % kinetic step
      if (step==n_steps_sample)
         % apply only 1/2 kinetic step when wavefunction is sampled
         U= u_kin_half;
      else
         U= u_kin;
      end
      psi_1=    ifft(U.*fft(psi_1));

      % Add contribution due to field interactions
      W1= E(tot_steps+1)*mu12;
      psi_1=    psi_1    + i*dt*(W1.*psi_0);
      t=t+dt;
      tot_steps=tot_steps+1;
      
   end
   n_sample=n_sample+1;
   pop= dot(psi_1,psi_1);
   pop_1(n_sample)= pop;
   
   %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
   % generate movie of wigner distributions%
   %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
   
  % W=Wig(psi_1);	
  W=wig2(conj(psi_1)',r); 
  figure(2);
   surf(W);
   shading interp;
   view(0,80);
  % axis ([0 140 0 140 -5 5]) ;     %axis([0 128 0 128 -0.05 0.05]);
  axis ([0 250 0 300 -5 5]) ; 
   caxis(.5*[(1.0e-008)*-.5  (1.0e-006)*0.4]);%-------------------------------------------------------
  % M(:,n_sample)=getframe;%-------------------------------------------------
   
   %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
   %%%% calculate and store psi_1 center (first normalize)
   %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
   psi_1_sqw=sum(psi_1.*conj(psi_1)); normh=(psi_1_sqw*dr); psi_1n=psi_1/sqrt(normh);
   
   loc(n_sample)=dot(r,(psi_1n.*conj(psi_1n)))*dr;
   time(n_sample)=conv_t*t;
   
   %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
   if (pop)
      r_1(n_sample)= real(dot(psi_1, (r.').*psi_1)/pop_1(n_sample));
      r_1_sqr(n_sample)= real(dot(psi_1, ((r.^2).').*psi_1)/pop_1(n_sample));
   end
   if (plot_wave & mod(tot_steps/n_steps_sample, update_freq)==0)
      set(psi_plot(2), 'YData', 10*wave_scale*abs(psi_1)+ ...
                                   (E_init+w0_1)*conv_en);
      set(psi_plot(3), 'YData', 10*wave_scale*real(psi_1)+ ...
                                   (E_init+w0_1)*conv_en);
      s= sprintf('\\psi_1: Time= %.0f fs', conv_t*t);
      title(s);
      drawnow;
   end
end  %end of "while" loop

p_final= pop_1(length(pop_1));


cpu_end= cputime;
cpu_elapsed= cpu_end-cpu_start
time_elapsed= etime( clock, clock_start) 

%%%%%%%%%%%%%%%%%%
% show movie!
%%%%%%%%%%%%%%%%%%
%movie(,1);
%movie2avi(M,'i2_chrp_mv2_long.avi','compression','cinepak','fps',3);%----------------------------------

figure(3)
plot(time,loc);
title('location of w.f. center');
xlabel('time (fs)');
ylabel('interatomic distance (Angstroms)');

%print -djpeg -r72 'i2_chrp_mv3.jpeg';  %------------------------------------------- 
return;