function wavefuncopt(optFileName, initFileName)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%  wavefuncopt
%  Optimizes a two level system governed by following Hamiltonian:
%
%          [ H_g        (mu*E(t)) ]
%      H = [                      ]
%          [ (mu*E(t))'     H_e   ]
%      
%         where H_g and H_e are the ground and excited state hamiltonians 
%         respectively:
%               H_{g/e}= -(d/dx)^2 +V_{g/e}(x).                
%         mu is the dipole operator and E(t) is the control field,
%    to maximize the objective J=Re{<\psi_T|P>}, 
%         where P is some target state.
%     
% 
%  Input Parameters (optional):
%    1) optFileName: Name of .mat file to save 
%                           optimized control and objective.
%    2) initFileName: Name of .mat file from which to initialize control.
%   
%    Both parameters are optional:
%    wavefuncopt 
%         - initializes from code and writes to default file: optFile.mat
%    wavefuncopt(optFileName)
%         - initializes from code and writes to optFileName.
%    wavefuncopt(optFileName, initFileName)
%         - initializes from initFileName and writes to optFileName.
%
%
%  Run can be halted by pressing the keyboard button. Command is then 
%    passed to the command line, and the workspace of the function can be 
%    directly viewed (and altered!) by typing 'dbup'. 
%  To resume run type 'return';
%
%  Comments:
%    0) Units: 
%       Harmonic Oscilator units. 
%    1) In the current code V_{g/e}(x) (V1 and V2 in code) 
%       are taken to be 
%          a) Displaced harmonic oscilators. (commented out)
%          b) V1=V_g a double well V_g(x)= a x^4 - b x^3
%             V2=V_e a harmonic oscilator.   (default)
%    2) For the propagation step a Split Operator Method is implemented:
%        psi(t+dt)=exp(iVdt/2)IFFT exp(ip^2 dt)FFT exp(iVdt/2) psi(t)
%        This is performed within the function evolve().
%        Note how exp(iVdt/2) psi is calculated by analytical 
%        diagonilzation of V. 
%        for details see Ch. 11.5.2
%
%    3) For details on the Global optimization method see Ch. 16.2
%       See also  S.E. Sklarz & D.J. Tannor, PRA 66 053619 (2002).
%
%    4) Two different objectives have been implemented:
%       a) J=Re{<\psi_T|P>}   (default)
%       b) J=|<\psi_T|P>|^2   (Commented out)
%       To switch objective 1) Exchange definition of objective 
%       in Objective().     2) Exchange comment marks on initialization
%       of chi(T). (Lines preceeding backwards propagation) 
%
%  Written by Shlomo Sklarz.   
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Everything in HO units

global P Nx dx V1 V2 E lambda dt;


tinitial=0;
tfinal=2*pi;                               %Length of time interval
tau=tfinal;                                 %Trap displacement function period

% Parameters of the calculation
xmax=12;                  % range of x dimension
Nx=2^6;                   % no of points
dx=2*xmax/Nx;
dt=.0005;                 % approx time step in each iteration
T=tfinal-tinitial;        % the time for the whole process
Nt=fix(T/dt);             % no of time steps
dt=T/Nt;                  % time step in each iteration 
Nplt=200;                 % No of timesteps per plot.

d=3;
d0=d;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%5
% Prepare coordinates.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%x=[-Nx/2:Nx/2-1]'*dx;      % x range vector
x=[-(Nx-1)/2:(Nx-1)/2]'*dx;
t=[0:Nt-1]'*dt;   

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%5
% Load data (optField and objective) from file.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
if nargin<2,
   E=.1./cosh((t-T/4)./(T/16)).^2.*exp(i*20*t);
   iter0=0;
   J=[0 0];
else
    load(initFileName);
    iter0=size(J,1)-1;
end;
if nargin<1, optFileName='optFile.mat'; end;

DE=nan*ones(Nt,1);
Ham=zeros(Nt,3);
%del=spdiags(ones(Nx,1)*[1 -2 1]/dx^2/2,[-1 0 1],Nx,Nx);
      dp=2*pi/(dx*Nx);
      pshft=dp*[-Nx/2:Nx/2-1]';
      p=fftshift(pshft);
      p2=p.^2/2;
mu=-1;
dinit=d;

midx=1; % movie frame index;
hfig=figure; subplot(311); subplot(312);subplot(325);subplot(326);
haxE=subplot(312);
hpltfld=plot(t,real(E),t,real(DE));
   set(hpltfld(2),'linewidth',2);
   title('Control');
   ylabel('E(t)');
   xlabel('time');
   axis([0,T,-8, 8]);
uicontrol('style','pushbutton','position', [430 10 60 20], ...
          'string','keyboard','callback','keyboard','userdata',1) 
drawnow;

%%%%%%%%%%%%%%%%%%%%%%%%%%%
%  Ground and excited state potentials 
%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Note: these can be changed at will.

% V1=(x.^2)/2; V2=((x-d).^2)/2;
V2=((x).^2)/2;  V1=.0625*(V2.^2-40*V2)+25; 
V12=1;   % Energy difference for graphical purposes only!

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Prepare initial wave function.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%psi0=exp(-(x).^2/2);  
psi10=exp(-(x+6.3246).^2*2.5/2); 
psi20=.1*psi10;
norm=sqrt((psi10'*psi10+psi20'*psi20)*dx);
psi10=psi10/norm;
psi20=psi20/norm;

haxWF=subplot(311);
hplt=plot(x,abs(psi10).^2,'g',x,V12+abs(psi20).^2,'g', ...
          x,.025*V1,x,V12-.2+.025*(V2),'b',x,x,'r',x,x,'r');
set(hplt(5:6),'visible','off');
ylabel('|\Psi(t)|^2');
axis([-10 10 0 1.6]);
drawnow;   


%%%%%%%%%%%%%%%%%%%%%%%%%%
% optimization parameters
%%%%%%%%%%%%%%%%%%%%%%%%%%

Niter=1000;
lambda=.01+100*(exp(-12*t)+exp(-(T-t)*12));
%P=[0*x; psi10]; % target state
P=[psi10(Nx:-1:1); 0*x]; % target state
%J=[J zeros(1,Niter)];  % 

psi1=psi10; psi2=psi20;
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%5
disp('Initial forwards propagation');
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
axes(haxWF); title('Forward propagation');
ylabel('|\Psi(t)|^2');

for idx=1:Nt;

   [psi1, psi2]= evolve(dt,mu*E(idx),psi1,psi2);
   
   if mod(idx,Nplt)==0
     set(hplt(1),'YData', abs(psi1).^2)
     set(hplt(2),'YData', V12+abs(psi2).^2);

     drawnow;
   end;
end;
     


J(iter0+1,:)=Objective([psi1; psi2]);
disp(['  J=',num2str(J(iter0+1,1))]);
disp(['      norm check: norm=' num2str((psi1'*psi1+psi2'*psi2)*dx)]);
subplot(325);plot([0:iter0]',J(1:iter0+1,:),'.-'); 
title('Objective');
xlabel('Iteration number');
ylabel('Jbar (-blue), J (-green)');

% M(midx)=getframe(hfig);midx=1;
% save('movFile0', 'M');

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%55
disp('Main optimization loop');
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

for iter=(iter0+1):(iter0+Niter)

  disp(['Iteration ' num2str(iter) ':']);



  % For Objective J=|<psi|P>|^2 uncomment the folowing lines: 
  %  chi1=exp(i*angle(P(1:Nx)'*psi1*dx))*P(1:Nx);
  %  chi2=exp(i*angle(P(Nx+1:2*Nx)'*psi2*dx))*P(Nx+1:2*Nx);

  % For Objective J=Real{<psi|P>} uncomment the folowing lines:
  chi1=P(1:Nx);
  chi2=P(Nx+1:2*Nx);
  %%%%%%%%%%%%%%%%%%%%%%%%%
  disp('  backwards propagation');
  %%%%%%%%%%%%%%%%%%%%%%%%%
  axes(haxWF); title('Backward propagation'); ylabel('|X(t)|^2');
  for idx=Nt:-1:1;

    [chi1, chi2]= evolve(-dt,mu*E(idx),chi1,chi2);
 
     if mod(idx,Nplt)==0
        set(hplt(1),'YData', abs(chi1).^2,'color','r')
        set(hplt(2),'YData', V12+abs(chi2).^2,'color','r');

        drawnow;  
        % M(midx)=getframe(hfig);midx=midx+1;
     end;
  end;  

  psi1=psi10; psi2=psi20; 
  %%%%%%%%%%%%%%%%%%%%%%%%%%
  disp('  forwards propagation');
  %%%%%%%%%%%%%%%%%%%%%%%%%
  axes(haxWF); 
    title('Forward propagation'); 
    ylabel('|\Psi(t)|^2- green, |X(t)|^2- red');
  set(hplt(5:6),'visible','on');

  for idx=1:Nt 
  
      % calculate new d(idx);
        DE(idx)=-i*mu'/lambda(idx)*(chi2'*psi1-psi2'*chi1)*dx;


      Ham(idx,:)=[2*dx*imag(chi1'*(V1.*psi1+ifft(p2.*fft(psi1))) ...
                           +chi2'*(V2.*psi2+ifft(p2.*fft(psi2)))), ...    
                  2*dx*imag(chi1'*(mu*DE(idx)*psi2)+chi2'*(mu*DE(idx))'*psi1),...
                  -lambda(idx)*abs(DE(idx))^2 ];
 

      % propagate chi with old E(idx)
   
        [chi1, chi2]= evolve(dt,mu*E(idx),chi1,chi2);
      
      % propagate psi with new E(idx)
        E(idx)=DE(idx);       

        [psi1, psi2]= evolve(dt,mu*E(idx),psi1,psi2);
     

      if mod(idx,Nplt)==0
          set(hplt(1),'YData', abs(psi1).^2,'color','g')
          set(hplt(2),'YData', V12+abs(psi2).^2,'color','g');
          set(hplt(5),'YData', abs(chi1).^2)
          set(hplt(6),'YData', V12+abs(chi2).^2);
          set(hpltfld(2),'Ydata', DE);
 
          drawnow;
          % M(midx)=getframe(hfig);midx=midx+1;
      end;
 
   end;   
  set(hplt(5:6),'visible','off');


   DE=DE*nan; % make invisible;

   set(hpltfld(1),'Ydata',real(E))
   set(hpltfld(2),'Ydata',real(DE)); 

   J(iter+1,:)=Objective([psi1;psi2]);
   disp(['  J=',num2str(J(iter+1,1))]);
   disp(['      norm check: norm=' num2str((psi1'*psi1+psi2'*psi2)*dx)]);

   if (iter~=0 & J(iter+1,1)<J(iter,1))
     disp('Error: Decrease in objective.');
     keyboard;  
   end;
   subplot(325);plot([0:iter]',J(1:iter+1,:),'.-'); 
       title('Objective');
       xlabel('iteration number');
     ylabel('Jbar (-blue), J (-green)');
   subplot(326);plot(t,real(sum(Ham'))); 
       title('Optimization Hamiltonian');
       xlabel('time');
   save(optFileName, 'E','J');

end;

keyboard;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function J=Objective(psi)
   global P dx E lambda dt;
   % Uncomment one of the following choices of objective:
   %    (appropriate changes must be made in main function!!)
   % 1) 
   %J=abs(P'*psi*dx).^2-E'*(lambda.*E)*dt;
   % 2)
   J(2)=real(P'*psi*dx); J(1)=J(2)-0.5*abs(E'*(lambda.*E))*dt;
return;


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [psi1, psi2]=evolve(dt,Om, psi1,psi2)
   global V1 V2 Nx dx;
   persistent  expiV12x p expip2;
   if isempty(expiV12x)
      dp=2*pi/(dx*Nx);
      pshft=dp*[-Nx/2:Nx/2-1]';
      p=fftshift(pshft);
      expip2=exp(-i*p.^2/2*abs(dt)/2);
      expiV12x=exp(-i*(V1+V2)*abs(dt)/2);
      psi1tmp=psi1;
      psi2tmp=psi2;
   end
   
   sgndt=sign(dt); % determines direction of propagation

   D=sqrt((V2-V1).^2+abs(4*abs(Om).^2));
   cosD=cos(D*dt/2); sinDovD=sin(D*dt/2)./D;

   psi1=ifft(expip2.^sgndt.*fft(psi1));
   psi2=ifft(expip2.^sgndt.*fft(psi2));

   psi1tmp=expiV12x.^sgndt.*((cosD+i*(V2-V1).*sinDovD).*psi1 ...
                     -2i*Om.*sinDovD.*psi2);
   psi2tmp=expiV12x.^sgndt.*((cosD+i*(V1-V2).*sinDovD).*psi2 ...
                     -2i*Om'.*sinDovD.*psi1);

   psi1=ifft(expip2.^sgndt.*fft(psi1tmp));
   psi2=ifft(expip2.^sgndt.*fft(psi2tmp));

return;

