% Plot WKB approximation to reflection of an obliquely upward-propagating
% light wave phi(x,z,t) = Re f(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 m1 >>1. In practice we take m1 = 2*pi
%
% Plot is for t = 0

%------User-specified parameters------------------------
  c0 = 1; % Light speed in vacuum
  m1 = 2*pi;% Incoming vertical wavenumber
  k = m1; % Hor. wavenumber
  n1 = 2; % Refractive index below gradient layer.
  n2 = 1; % Refractive index above gradient layer (n2 > n1).
  f1 = 1; % Incoming wave amplitude
  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;
  ximatch = -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;
  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 + ximatch/mu13;
  kmatch = floor((zmatch - zmin)/dz) + 1; % gridpoint index below turning point
  
  % Find mint = int_zmin^z m(zeta) dzeta at each gridpoint below
  % turning point using
  % trapezoidal rule. mint(k) = dz*(0.5*m(1)+m(2)+...+m(k-1)+0.5*m(k))
  
  mint = 0.5*dz*[0 cumsum(m(1:kt-1)+m(2:kt))];
  alpha = mint(kt) + m(kt)*0.5*remt*dz; % Get phase integral up to zt using.
    % trapezoidal rule, noting m(kt+remt) = 0: 
    % alpha = dz*(0.5*m(1) + m(2) + ....m(kt-1) +(0.5+remt)*m(kt))
  r = exp(1i*(2*alpha - pi/2));
  
  zWKB = z(1:kmatch); % Use WKB solution up to matching height zm
  zAiry = z(kmatch+1:end); % Use Airy solution above the matching
                           % height
  f = zeros(1,nz);		   
  f(1:kmatch) = f1*sqrt(m1./m(1:kmatch)).*exp(1i*mint(1:kmatch));
  f(1:kmatch) = f(1:kmatch) + r*conj(f(1:kmatch)); % Add reflected wave
  c = -f1*2*1i*sqrt(pi*m1)/mu13 * exp(1i*(alpha+pi/4));
  f(kmatch+1:nz) = c*airy(mu13*(z(kmatch+1:nz) - zt));
  [X F] = meshgrid(x,f);
  PHI = real(exp(1i*k*X).*F);

  % Contour plot
  
  contourf(x,z,PHI)
  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')
  hold off