% parabolic_surface_refract.m
%
% Trace incoming parallel rays for parabolic-surface "lens"
%   defined by f(y,z) = z-zc - a*y^2 = 0.
%
% Represent rays by a position y at z, and also a direction unit vector
%   rhat = [ry; rz]; z is the direction along the optical axis.
%
% Use coordinate-free form of Snell's Law in terms of surface normal vector.

global gY gRhat gZ0 gA gZc

a = 0.5;       % parabolic parameter
zc = -0.5;     % shift of parabola center
dy = 0.06;     % increment of initial y position of rays
ymax = 6*a;    % max y ray to trace
zmin = -6*a;   % launch z
zmax =  6*a;   % max z
n1 = 1;        % index to left of parabola
n2 = 2;        % index to right of parabola

% In the following function we need to pass extra parameters, 
%   hence the global statement.
%     gY -> y at z0
%     gRhat -> ray direction (unit) vector, in form [ry; rz]
%     gZ0 -> reference coordinate z0
%     gA -> copy of A
%     gZc -> copy of zc
% Note: global variables are labeled with a 'g' to avoid confusion
gA = a;
gZc = zc;

% function for finding intersections with the parabolic surface
function fout = surfacefunc(z)
  % z is the position along the optical axis
  global gY gRhat gZ0 gA gZc

  % extrapolate ray straight to z
  %   y = y0 + (z-z0)*ry/rz;
  ray_y = gY + (z-gZ0)*gRhat(1)/gRhat(2);

  % function to define surface via f(y,z) = (z-zc) - a*y^2 = 0
  fout = (z-gZc) - gA*ray_y^2;
end %function

% function for finding the normal vector to the surface
%   returns gradient of f = z-zc - a*y^2
function fout_vec = gradfunc(y,z)
  global gA
  fout_vec = [ ...
    -2*gA*y;    ... % partial f/partial y
    1          ... % partial f/partial z
  ];
  fout_vec = fout_vec/sqrt(dot(fout_vec,fout_vec)); % make unit vector
end %function


% plot parabola, using parametric form
dt = 0.01;
t=(zmin:dt:zmax)';
plot(a*t.^2+zc, t,'b-'); % plot as blue line

% set axes using z dimensions, and square aspect ratio
axis([zmin, zmax, zmin, zmax], "square");
xlabel('z (mm)');
ylabel('y (mm)');
title(sprintf('ray trace of parabolic interface, n = %.2f',n2));
hold on; % subsequent plots will be overlayed

%%% loop over rays
nrays = 2*floor(ymax/dy);
for ray = 0:nrays,
  y0 = (ray-nrays/2) *dy;
  yarr = [y0];   % array of y positions to plot
  zarr = [zmin]; % array of z positions

  % find first intersection with the parabolic surface
  %   look for intersection (root) between zmin and 10*zmax
  rhat = [0; 1]; % direction of initial ray, horizontal = zhat
  gY = y0; gRhat = rhat; gZ0 = zmin;
  z = fzero('surfacefunc', [zmin; 10*zmax]);
  y = gY;

  % add point to plot
  yarr = [yarr; y]; zarr = [zarr; z];

  % compute refracted ray direction, update direction vector rhat
  nhat = gradfunc(y,z); % surface normal vector
  nhat3 = [0; nhat]; rhat3 = [0; rhat];    % add x dimension for cross product
  crossthing = (n1/n2)*cross(nhat3,rhat3); % (n1/n2)*(n x r)
  rhat = (n1/n2)*rhat ...
         + ( sign(dot(nhat,rhat))*sqrt(1-dot(crossthing,crossthing)) ...
             - (n1/n2)*dot(nhat,rhat) )*nhat;

  % extrapolate to the end of the optical axis
  z0 = z; z = zmax;
  yarr = [yarr; y + (z-z0)*rhat(1)/rhat(2)]; zarr = [zarr; z];

  % plot ray
  plot(zarr, yarr, 'r-'); % plot as red line
end %for

feff = n2/(2*a*(n2-n1)); % effective, paraxial focal length
% mark effective focal length, account for offset of surface
plot([feff+zc], [0], 'g.');
hold off;