#!/usr/bin/env Python3 import numpy as np from pdb import set_trace import matplotlib.pyplot as plt from matplotlib.ticker import FormatStrFormatter from matplotlib.colors import LogNorm import time import sys sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics') from fourierOptics_new import * if __name__ == '__main__': plt.rc('font', family='serif') plt.rc('axes', labelsize='large') # Start time counter start = time.time() # Pick a wavelength l = 30e-3 # Initiate Fourier Optics fou = FourierOptics(l=l) # Pick sampling parameters dpix = 1/50 n, m = 2001, 2001 print ('Your resolution is {}x{} with a pixel size of {}'.format(n,m,dpix)) # Initiate the profiles class and pick a profile (tophat, obscured tophat or spider) prof = Profiles() A = prof.top_hat(2.5, n, m, dpix=dpix, oversample=1) # Create an input grid x, y = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m)) # Alternatively, create a gaussian profile # gauss = GaussianBeam(0.3, l*1000) # A = gauss.fundamental_mode(0,n,m,np.max(x)) # Create a lens with a focal length ~100 mm, with a clear aperture radius of 4 mm and thickness of 1 mm. L1 = LensObject('N-BK7', 1, 4, name='L1') L1.set_first_surface('EVEN ASPHERE', R=11.65, k=-1.10, a=[0.0, 3.69067205091E-5, -1.285461192617E-8, -1.400167659754E-10, -2.51311664812E-13, 5.01789877647E-16, 5.855871513042E-18, -1.127794407419E-20]) L1.set_second_surface('SPHERICAL', R=1e99) print('Focal length of lens {} is {:.1f} mm.'.format(L1.name, L1.f)) # set_trace() # Calculate the phase delay introduced by the lens t_l = L1.phase_delay(x,y,l, paraxial=False) U_in = A * t_l # Performing the propagation using Fresnel approximation to the focal plane. oC = [x.min(), x.max(), y.min(), y.max()] savename = False Exyz_og,_,_ = fou.fresnel_propagation(U_in, x, y, dz=L1.f, show=True, save=savename, outputCoordinates=oC) # Add PSD wavefront error to the lens wfe = WavefrontError() h_data = wfe.generate_height_map(1e-2, 2, 1, n, m) L1.add_wavefront_error(h_data) t_l = L1.phase_delay(x,y,l, paraxial=False) # The input distribution gets a curved wavefront after the lens U_in = A * t_l # Performing the propagation using Fresnel approximation to the focal plane. Exyz_psd,_,_ = fou.fresnel_propagation(U_in, x, y, dz=L1.f, show=False, save=savename, outputCoordinates=oC) # Now we try a zernike x_z, y_z = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, 500), np.linspace(-m/2*dpix, m/2*dpix, 500)) zern = wfe.zernike_cartesian([0, 0, 0, 1, 0, 0, 0, 1], x_z, y_z)/75 zern2 = np.pad(zern, ((750,751),(750,751)), 'constant', constant_values=(0)) zern = zern2 * prof.top_hat(4, n,m, dpix) L1 = LensObject('N-BK7', 1, 4, name='L1') L1.set_first_surface('EVEN ASPHERE', R=11.65, k=-1.10, a=[0.0, 3.69067205091E-5, -1.285461192617E-8, -1.400167659754E-10, -2.51311664812E-13, 5.01789877647E-16, 5.855871513042E-18, -1.127794407419E-20]) L1.set_second_surface('SPHERICAL', R=1e99) L1.add_wavefront_error(zern) t_l = L1.phase_delay(x,y,l, paraxial=False) U_in = A * t_l Exyz_zern,_,_ = fou.fresnel_propagation(U_in, x, y, dz=L1.f, show=True, save=savename, outputCoordinates=oC) E1, fx1, fy1 = fou.modulation_transfer_function(np.abs(Exyz_og)**2, x,y,dpix) E2, fx2, fy2 = fou.modulation_transfer_function(np.abs(Exyz_zern)**2, x,y,dpix) plt.figure() plt.plot(fx1[1000:], E1[1000, 1000:],'b-', label='Diffraction limited') plt.plot(fx2[1000:], E2[1000, 1000:],'r-', label='Astigmatism') plt.xlim(0,1.67) plt.ylim(0,1) plt.xlabel('Spatial frequency (cycles/aperture)') plt.ylabel('MTF') plt.legend() plt.grid() plt.show() generate_some_images = False if generate_some_images: # Disturbed PSF plt.figure() plt.imshow(np.abs(Exyz_psd), extent=[x.min(), x.max(), y.min(), y.max()], cmap='hot') plt.xlim(-10,10) plt.ylim(-10,10) plt.colorbar() plt.savefig('./figures/PSF_with_psd_error.png', bbox_inches='tight', transparent=True) plt.close() # Wavefront error plt.figure() plt.imshow(h_data, cmap='jet') plt.colorbar() plt.savefig('./figures/psd_error.png', bbox_inches='tight', transparent=True) plt.close() # Original PSF plt.figure() plt.imshow(np.abs(Exyz_og), extent=[x.min(), x.max(), y.min(), y.max()], cmap='hot') plt.xlim(-10,10) plt.ylim(-10,10) plt.colorbar() plt.savefig('./figures/PSF_no_psd_error.png', bbox_inches='tight', transparent=True) plt.close() # Zernike PSF (some coma and also astigmatism) plt.figure() plt.imshow(np.abs(Exyz_zern), extent=[x.min(), x.max(), y.min(), y.max()], cmap='hot') plt.xlim(-10,10) plt.ylim(-10,10) plt.colorbar() plt.savefig('./figures/PSF_with_zernike_error.png', bbox_inches='tight', transparent=True) plt.close() # Wavefront error, zernike plt.figure() plt.imshow(zern, cmap='jet') plt.xlim(500,1500) plt.ylim(500,1500) plt.colorbar() plt.savefig('./figures/zernike_error.png', bbox_inches='tight', transparent=True) plt.close() # print(np.sum(np.abs(Exyz)**2)) # set_trace() # Plot a cross section of the new field (useful to better see the diffraction rings) # fig = plt.figure() # ax = fig.add_subplot(111) # ax.grid(True, alpha=0.3) # ax.semilogy(np.linspace(np.min(oC), np.max(oC),n), np.abs(Exyz[1000,:]**2), 'r-') # ax.set_xlabel('x (mm)') # ax.set_ylabel('Intensity') # ax.tick_params(axis ='both', direction='in', which='both') # plt.savefig('simple_lens_slice.pdf', bbox_inches='tight', transparent=True) # plt.show() # The following produces a cross-section of the propagation for various z, so that you can see how the beam evolves. if False: Ics = [] Ics_x = [] Isum = [] xmax = [] zList = np.linspace(20,220,20).tolist() xin, yin = x, y for i in zList: title = f"{zList.index(i)+1:03d}" print (' Calculating {0}/{1}'.format(title, f"{len(zList):03d}"), end="\r") # Propagate the wavefront # Exyz, x, y = fou.planeWavePropagationFFTW(A, dz=i*l, dpix=dpix, oversample=1, show=False, save=False)#, scaleInputArray=(n,m)) Exyz, x, y = fou.fresnel_propagation(U_in, xin, yin, dz=i, show=False, save=False, outputCoordinates=oC) # Process the results by taking a slice of the data Isum.append(np.sum(np.abs(Exyz)**2)) Ics.append((np.abs(Exyz[int(Exyz.shape[0]/2)])**2)) Ics_x.append(x.real[0]) y_ext = [y.real[0,0],y.real[-1,0]] # Reset the data x = None y = None # Plotting all slices fig = plt.figure(figsize=(8,4)) ax = fig.add_subplot(311) ax.tick_params(axis ='both', direction='in') # Plotting the square root of the amplitude for better contrast im = ax.imshow(np.sqrt(np.asarray(Ics)).T, cmap='hot', aspect='auto', interpolation='bicubic', extent=[zList[0],zList[-1],y_ext[0],y_ext[1]]) ax.yaxis.set_major_formatter(FormatStrFormatter('%.1E')) # Plotting the intensity relative to the input ax2 = fig.add_subplot(312) ax2.plot(zList, np.asarray(Ics)[:,int(np.asarray(Ics).shape[1]/2)], 'k-') ax2.yaxis.set_major_formatter(FormatStrFormatter('%.1E')) ax2.set_xlabel('Propagation distance (mm)') ax2.tick_params(axis ='both', direction='in') ax2.set_ylabel('Amplitude') ax2.set_xlim([np.min(zList), np.max(zList)]) ax2.grid(True, alpha=0.3) ax3 = fig.add_subplot(313) ax3.grid(True, alpha=0.3) ax3.set_xlabel('Propagation distance (mm)') ax3.tick_params(axis ='both', direction='in') ax3.set_ylabel('Normalized power') ax3.set_xlim([np.min(zList), np.max(zList)]) ax3.plot(zList, np.asarray(Isum)) # Saving the figure? Then show. plt.tight_layout() # plt.savefig('fresnel_cross.pdf', bbox_inches='tight') # plt.savefig('fresnel_cross.png', bbox_inches='tight') plt.savefig('simple_lens_cross.pdf', bbox_inches='tight', transparent=True) plt.show() # Closing down the script and stopping the timer. end = time.time() print('The program took {:.3f} seconds!'.format((end-start)))