%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
% Class:     Psych 221/EE 362
% File:      WaveAberrationPSF
% Author:    Patrick Maeda
% Purpose:   Calculate and Plot PSF of Wave Aberration Function
% Date:      03.08.03	
%	
% Matlab 6.1:  03.09.03
% 
% revised by Al Ahumada, 25 Jun 08
%  Matrix operations for increased speed.
%  Clarifies expansion.
%  Allows different image sizes in number and width,
%  allowing a larger range of pupil sizes.
%  Section plots include lines at first zero of 
%  aberration-free point spread.
%  Aberration-free calculations optional (aberration0).
%  Plotting optional (plotting).
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
% This file calculates and plots the PSF of the Wave Aberration function specified by:
% jmax = highest mode in Wave Aberration Expansion
% Wrmsj = rms Wave Aberration coefficients of modes 0 to jmax
% d = pupil diameter in mm
% lambda = wavelength in nm
%
% The Zernike Polynomial definitions used are derived from:
% Thibos, L., Applegate, R.A., Schweigerling, J.T., Webb, R., VSIA Standards Taskforce Members,
% "Standards for Reporting the Optical Aberrations of Eyes"
% OSA Trends in Optics and Photonics Vol. 35, Vision Science and its Applications,
% Lakshminarayanan,V. (ed) (Optical Society of America, Washington, DC 2000), pp: 232-244. 
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

clear all;

aberration0 = 1;
% aberration0 = 0;

plotting = 1;
% plotting = 0;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Wave Aberration definition
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

disp('Maximum Mode Number j')
jmax=14                  %[INPUT] highest mode number from single indexing scheme
nj=[0:jmax]; mj=nj ;
for j=0:jmax
   n=ceil((-3+sqrt(9+8*j))/2);   %highest power or order of the radial polynomial term
   nj(j+1)=n;   %highest power or order of the radial polynomial term
   mj(j+1)=2*j-n*(n+2);                %azimuthal frequency of the sinusoidal component
end
% nj =
% 
%      0     1     1     2     2     2     3     3     3     3     4     4     4     4     4
% mj =
% 
%      0    -1     1    -2     0     2    -3    -1     1     3    -4    -2     0     2     4
     
disp('Pupil Diameter (mm), RMS Wave Aberration Coefficients (micron), and Wavelength (nm)')
d=5.4;                   %[INPUT] pupil diameter in mm (3 to 8 mm)
d=3;                   %[INPUT] pupil diameter in mm (3 to 8 mm)
PupilDiameter=d
Wrmsj=[0 0 0 0.4164 0 0.135 0.074 -0.092 0.011...  %[INPUT] rms wavefront error coefficient in microns
       -0.12 -0.038 0.016 0.085 -0.06 0.047]'
Wrmsj=[0 0 0 0 0.2 0 0 0 0 0 0 0 0 0 0]'  %[INPUT] focus error only
      
% Wrmst=0;
% for j=0:jmax
%    Wrmst=Wrmst+Wrmsj(j+1)^2;
% end
% Wrmstotal=sqrt(Wrmst)    %total rms wavefront error in um
Wrmstotal=sqrt(sum(Wrmsj.*Wrmsj))    %total rms wavefront error in um

lambda=570               %[INPUT] wavelength in nm

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Convert units for calculation
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

Wrmsj=Wrmsj*1e-3;          %rms wavefront error coefficients in mm
Wrmstotal=Wrmstotal*1e-3;  %total rms wavefront error in mm
lambda=lambda*1e-6;        %wavelength in mm
dw=d/lambda;               %pupil diameter in number of wavelengths
PRw=0.5*dw;                %pupil radius in number of wavelengths
apw=pi*PRw^2;              %pupil area in wavelength^2

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Set-up x,y grid
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% xwmin=-25000;  %minimum x-coordinate in number of wavelengths
% xwmax=25000;   %maximum x-coordinate in number of wavelengths
% ywmin=-25000;  %minimum y-coordinate in number of wavelengths
% ywmax=25000;   %maximum y-coordinate in number of wavelengths
% dxw=250;       %x-coordinate pixel width in number of wavelengths
% dyw=250;       %y-coordinate pixel width in number of wavelengths
% 
% xw=xwmin:dxw:xwmax;   %x-coordinates in number of wavelengths
% yw=ywmin:dyw:ywmax;   %y-coordinates in number of wavelengths
% Imax=length(xw);
% Jmax=length(yw);
Imax = 512; 
Imax = 201; 
Jmax = Imax;
Imax2 = floor(Imax/2);
Jmax2 = floor(Jmax/2);
IembedFactor = 5; % ratio of image size to pupil diameter
 IembedFactor = 25000/PRw ;
JembedFactor = IembedFactor; % ratio of image size to pupil diameter
dxw = IembedFactor*PRw/(Imax2) ;
dyw = JembedFactor*PRw/(Jmax2) ;
xw = dxw*([1:Imax]-(1+Imax)/2);
yw = dyw*([1:Jmax]-(1+Jmax)/2);
% xwmin = xw(1) ; xwmax = xw(Imax);
% ywmin = yw(1) ; ywmax = yw(Imax);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Set-up circular pupil
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% for I=1:Imax    
%     for J=1:Jmax
%        P(I,J)=(sqrt(xw(I)^2+yw(J)^2) <= PRw);
%     end
% end

P=(repmat((xw.*xw)',[1 Jmax])+repmat(yw.*yw,[Imax 1])) <= (PRw*PRw);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Compute Wave Aberration function
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% W=zeros(length(xw),length(yw));
% for j=0:jmax
%    n=ceil((-3+sqrt(9+8*j))/2);   %highest power or order of the radial polynomial term
%    m=2*j-n*(n+2);                %azimuthal frequency of the sinusoidal component
%    W=W+Wrmsj(j+1)*zernike(n,m,xw,yw,dw);
% end

W=zeros(size(P));
for j1=1:jmax+1
    if Wrmsj(j1) ~= 0
        W=W+Wrmsj(j1)*zernike(nj(j1),mj(j1),xw,yw,dw);
    end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 
% Compute PSF
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if aberration0
    PSF0=fft2(P)/apw;
    PSF0=fftshift(PSF0);
    PSF0=PSF0.*conj(PSF0);
    PSF0=rot90(PSF0);
    PSF0=flipud(PSF0);
end
PSF=fft2((P.*exp(-i*2*pi*W/lambda)))/apw;
if plotting
PSF=fftshift(PSF);
PSF=PSF.*conj(PSF);
PSF=rot90(PSF);
PSF=flipud(PSF);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Plot PSF
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% umin=-0.002;    %minimum thetax-coordinate in radians
% umax=0.002;     %maximum thetax-coordinate in radians
% vmin=-0.002;    %minimum thetay-coordinate in radians
% vmax=0.002;     %maximum thetay-coordinate in radians
% du=0.00002;     %thetax-coordinate pixel width in radians
% dv=0.00002;     %thetay-coordinate pixel width in radians
% u=umin:du:umax;   %thetax-coordinates in radians
% v=vmin:dv:vmax;   %thetay-coordinates in radians
du=0.00002*(100/128)*(250/dxw);     %thetax-coordinate pixel width in radians
du=0.5/(Imax2*dxw);     %thetax-coordinate pixel width in radians
dv=0.5/(Jmax2*dyw);     %thetay-coordinate pixel width in radians

u = du*([1:Imax]-(1+Imax2));
v = dv*([1:Jmax]-(1+Jmax2));
umin = u(1) ; umax = u(Imax);
vmin = v(1) ; vmax = v(Jmax);

angfac = 1000 ; anglab = 'mrad';
angfac = 60*180/pi ; anglab = 'arc min';
plotn = aberration0 + 1;

figure
if aberration0
subplot(plotn,1,1)
scale=(2)^7/max(max(PSF0));
image(u*angfac,v*angfac,PSF0*scale)  %scaled for saturated display of image
%imagesc(u*angfac,v*angfac,PSF0)
axis image
xlabel(['\theta_{x} (' anglab ')'])
ylabel(['\theta_{y} (' anglab ')'])
axis square
axis xy
title(['PSF of Zero Aberration System, Pupil Diameter = ',...
        num2str(PupilDiameter), 'mm'], 'FontSize', 10);
colormap gray
end
%figure
subplot(plotn,1,plotn)
scale=(2)^7/max(max(PSF));
image(u*angfac,v*angfac,PSF*scale)   %scaled for saturated display of image
%imagesc(u*angfac,v*angfac,PSF)
axis image
axis xy
xlabel(['\theta_{x} (' anglab ')'])
ylabel(['\theta_{y} (' anglab ')'])
axis square
title(['PSF of Aberrated System, RMS Wavefront Error = ',num2str(Wrmstotal/lambda),'\lambda'],'FontSize', 10);
colormap gray

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 
% Plot PSF Cross-sections
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

figure
if aberration0
subplot(plotn,2,1)
plot(u*angfac, PSF0(floor((Imax+1)/2),:)/max(max(PSF0)) ...
    ,(angfac/dw)*[-1 -1],[0 1],(angfac/dw)*[+1 +1],[0 1])
xlabel(['\theta_{x} (' anglab ')'])
ylabel('Normalized Amplitude')
axis([umin*angfac umax*angfac 0 1])
axis square
title(['PSF of Zero Aberration System, PD = ',...
        num2str(PupilDiameter), 'mm'], 'FontSize', 10);
subplot(plotn,2,2)
plot(v*angfac, PSF0(:, floor(Jmax+1)/2)/max(max(PSF0)) ...
    ,(angfac/dw)*[-1 -1],[0 1],(angfac/dw)*[+1 +1],[0 1])
xlabel(['\theta_{y} (' anglab ')'])
ylabel('Normalized Amplitude')
axis([vmin*angfac vmax*angfac 0 1])
axis square
title(['PSF of Zero Aberration System, PD = ',...
        num2str(PupilDiameter), 'mm'], 'FontSize', 10);
end
subplot(plotn,2,1+2*(plotn-1))
plot(u*angfac, PSF(floor((Imax+1)/2),:)/max(max(PSF)) ...
    ,(angfac/dw)*[-1 -1],[0 1],(angfac/dw)*[+1 +1],[0 1])
xlabel(['\theta_{x} (' anglab ')'])
ylabel('Normalized Amplitude')
axis([umin*angfac umax*angfac 0 1])
axis square
title(['PSF of Aberrated System, Wrms = ', num2str(Wrmstotal/lambda),'\lambda'],'FontSize', 10);
subplot(plotn,2,2+2*(plotn-1))
plot(v*angfac, PSF(:, floor((Jmax+1)/2))/max(max(PSF)) ...
    ,(angfac/dw)*[-1 -1],[0 1],(angfac/dw)*[+1 +1],[0 1])
xlabel(['\theta_{y} (' anglab ')'])
ylabel('Normalized Amplitude')
axis([vmin*angfac vmax*angfac 0 1])
axis square
title(['PSF of Aberrated System, Wrms = ', num2str(Wrmstotal/lambda),'\lambda'],'FontSize', 10);
end