% Plot WKB approximation to bending 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) increases 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 = 1, N2 = 2.
% (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 = 1; % Refractive index below gradient layer.
  N2 = 2; % Refractive index above gradient layer (N2 > N1).
  zmin = -2;  % Min, max and increment in z for plots
  zmax = 2;
  dz = 0.05;
  xmin = 0;   % Min, max and increment in x for plots
  xmax = 4.4;
  dx = 0.1;

%------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 Mint = int_zmin^z M(zeta) dzeta at each gridpoint 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:nz-1)+M(2:nz))]; 
  y = sqrt(M1./M).*exp(1i*Mint);  % Incoming wave of unit amplitude
  [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)
  title('WKB approximation to refraction through 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.8*xmax, 1.7,['N_2 = ' num2str(N1)],'Color','g')
  text(0.8*xmax, 1.5,['M_2 = ' num2str(0.1*round(10*M(nz)))],'Color','g')
  text(0.65*xmax, 0,  ['N(z) = ' num2str(0.5*(N1+N2)) ' + '...
                      num2str(0.5*(N2-N1)) 'erf(z)'],'Color','g')
  hold off