% Plot WKB approximation to reflection of an obliquely upward-propagating
% light wave u(x,z,t) = Re y(z)*exp(i(Kx-om*t)) by a layer in
% which refractive index N(z) decreases with z.  Local dispersion relation is
% om^2 = c0^2*(K^2 + M^2(z))/n^2(z).

% Assumptions:
% (1) N(z) = N1 + (N2-N1)*(1+erf(z))/2  where z is nondimensional height.
%     Thus N -> N1 as z->-inf, N -> N2 as z->inf.  Here N1 = 2, N2 = 1.
% (2) Nondimensionalized light speed in vacuum c0 = 1
% (3) Incoming wave of unit amplitude is at 45 degrees to the vertical. Hence
%     incoming vertical wavenumber M1 = hor. wavenumber K.
% (4) For WKB to hold. need eps^-1 = K >>1. In practice we take K = 2*pi
%
%  Translation table to m(z)/epsilon notation in lecture notes:
%    M(z) = K*m(z) = m(z)/eps
%
% Plot is for t = 0

%------User-specified parameters------------------------
  c0 = 1; % Light speed in vacuum

  K = 2*pi;% Incoming horizontal wavenumber (= 1/eps)
  M1 = K; % Incoming vertical wavenumber (M1=K gives incidence angle of 45 deg) 
  N1 = 2; % Refractive index below gradient layer.
  N2 = 1; % Refractive index above gradient layer (N2 > N1).
  zmin = -3;  % Min, max and increment in z for plots
  zmax = 1;
  dz = 0.05;
  xmin = 0;   % Min, max and increment in x for plots
  xmax = 4.4;
  dx = 0.1;
  wmatch = -2; % Transition from Airy to WKB solution

%------Calculations--------------------------------------
  om = sqrt(c0^2*(K^2 + M1^2)/N1^2); % Wave frequency
  x = xmin:dx:xmax;
  z = zmin:dz:zmax;
  nz = length(z);
  N = N1 + (N2-N1)*(1+erf(z))/2;
  M2 = N.^2*om^2/c0^2 - K^2;
  M = sqrt(N.^2*om^2/c0^2 - K^2);
  
  % Find turning point zt where M=0 (and N = c*K/om = sqrt(2)).
  
  nt = c0*K/om;
  zt = interp1(N,z,nt,'spline');
  kt = floor((zt - zmin)/dz) + 1; % gridpoint index below turning point
  remt = (zt - z(kt))/dz;				  
  mu = -(M(kt+1)^2 - M(kt)^2)/dz; % mu = -dM^2/dz(zt)
  mu13 = mu^(1/3);
  zmatch = zt + wmatch/mu13;
  kmatch = floor((zmatch - zmin)/dz) + 1; % gridpoint index below turning point
  
  % Find phase integral Mint = int_zt^z M(zeta) dzeta at each gridpoint below
  % turning point using the trapezoidal rule:
  %     Mint(k) = dz*(0.5*M(1)+M(2)+...+M(k-1)+0.5*M(k))-alpha
  % where (noting M(kt+remt) = 0)
  %     alpha = dz*(0.5*M(1) + M(2) + ....M(kt-1) +(0.5+remt)*M(kt))
  Mint = 0.5*dz*[0 cumsum(M(1:kt-1)+M(2:kt))];
  alpha = Mint(kt) + M(kt)*0.5*remt*dz; 
  Mint = Mint - alpha;

  r = -1i;  % WKB reflection coefficient
  
  zWKB = z(1:kmatch); % Use WKB solution up to matching height zm
  zAiry = z(kmatch+1:end); % Use Airy solution above the matching
                           % height
  y = zeros(1,nz);		   
  y(1:kmatch) = sqrt(M1./M(1:kmatch)).*exp(1i*Mint(1:kmatch)); % Incoming wave
  y(1:kmatch) = y(1:kmatch) + r*conj(y(1:kmatch)); % Add reflected wave
  b = 2*sqrt(pi*M1/mu13) * exp(-1i*pi/4);
  y(kmatch+1:nz) = b*airy(mu13*(z(kmatch+1:nz) - zt));
  [X Y] = meshgrid(x,y);
  u = real(exp(1i*K*X).*Y);

  % Contour plot
  
  contourf(x,z,u)
  colormap gray
  colorbar
  axis equal
  xlabel('x')
  ylabel('z')
  hold on
  plot([xmin xmax],[-1 -1],'b--',[xmin xmax],[1 1],'b--' ,'LineWidth',2)
  plot([xmin xmax],[zt zt],'r-',...
       [xmin xmax],[zmatch zmatch],'r-.','LineWidth',2)
  title('WKB approximation to reflection at a density gradient')
  
  text(0.8*xmax,-1.5,['N_1 = ' num2str(N1)],'Color','g')
  text(0.8*xmax,-1.65,['K = \epsilon^{-1} = 2\pi'],'Color','g')
  text(0.8*xmax,-1.85,['M_1 = 2\pi'],'Color','g')
  text(0.65*xmax, 0.5,  ['N(z) = ' num2str(0.5*(N1+N2)) ' + '...
                      num2str(0.5*(N2-N1)) 'erf(z)'],'Color','g')
  text(0.05*xmax, -0.4,'z_t-\delta','Color','r')
  text(0.05*xmax, 0.3,'z_t','Color','r')
  hold off
 