% ThruFocusPSF.m: Calculate the thru-focus PSF
%
close all; clear; clc

%Make the array 5 times the size of the pupil, 
% for adequate PSF resolution
Asize = 5;
N = 101; %Size of array;
x = linspace(-Asize,Asize,N);
[X,Y] = meshgrid(x,x);
R = sqrt(X.^2 + Y.^2);
%R is now an array with values equal to the distance
% from the center of the N x N array:

%Create a circular, uniform pupil:
P = zeros(N,N);
P(R<=1) = 1;

%% Study defocus effect
%The defocus is in unit of waves at the edge of the pupil:
Defocus = -2:.5:2; %Array of defocus amounts to use on Pupil
NDefocus = length(Defocus);

%%
% Calculate the PSF through focus:
PSFs = zeros(N,N,1,NDefocus);
for ii = 1:NDefocus
   DF = exp(i*2*pi*Defocus(ii)*R.^2);
   DF(P==0) = 0; %Truncate to aperture
   DP = P.*DF; %Defocus phase added to pupil
   %Transform and square the pupil to get the PSF array:
   psf = abs(fftshift(fft2(ifftshift(DP)))).^2;
   %Normalize the psf:
   psf = psf/max(max(psf));
   PSFs(:,:,1,ii) = psf;
   %Calculate the MTFs
   %We don't use 'fftshift', because we only want to extract
   % a central slice, and that will be at the edge of the
   % array before shifting.
   mtf = abs(fft2(psf));
   mtf = mtf(1,1:20)';
   mtf = mtf/max(mtf); %Normalize to 1
   MTFs(:,ii) = mtf;
end

%% Create a montage of the through-focus PSFs
figure(1)
montage(PSFs)
title('-2\lambda : 0.5\lambda : 2\lambda  Misfocus')

%Plot the associated MTFs:
xm = linspace(0,1,length(mtf)) ;
figure(2)
for ii = 1:NDefocus
   subplot(3,3,ii)
   plot(xm,MTFs(:,ii))
   title(['Defocus: ' num2str(Defocus(ii)) ' \lambda']);
end