#!/usr/bin/env Python3 import numpy as np from pdb import set_trace import json import matplotlib.pyplot as plt from matplotlib.ticker import FormatStrFormatter from matplotlib.colors import LogNorm import sys sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics') from fourierOptics import * def saveVariableToTXT(variable, filename): variableFile = open(filename, 'w') if type(variable) == np.ndarray: json.dump(variable.tolist(), variableFile) elif type(variable) in [list, str, float]: json.dump(variable, variableFile) else: print ('Not a list, array or string. Check variable type.') variableFile.close() return # Load a variable from a json text file. def loadVariableFromTXT(filename): variableFile = open(filename, 'r') variable = json.load(variableFile) variableFile.close() return variable def scale_zernike_coefficients(c, rms): # Scaling the coefficients to some total value rms_coeffic = np.sqrt(np.sum( np.asarray(c)**2)) scale_factor = rms / rms_coeffic c = (scale_factor * np.asarray(c)).tolist() rms_coeffic = np.sqrt(np.sum( np.asarray(c)**2)) return c def expected_strehl(rms, l): return np.exp( -(2*np.pi*rms/l)**2 ) if __name__ == '__main__': plt.rc('font', family='serif') plt.rc('axes', labelsize='large') # Parameters dpix = 160 n, m = 1001, 1001 oC = [-.5, .5, -.5, .5] l = 1.35e-3 # wavelength zern_rms = 45e-5 # RMS of the zernike component of the WFE midspat_rms = 21.8e-5 # RMS of the PSD component of the WFE rmsList = np.linspace(1e-4, 1e-3, 25) # 1 to 1000 nm WFE RMS in mm number_of_iterations = 50 # Start of the ELT-lens system # Initiate Fourier Optics, profile and wavefront error objects fou = FourierOptics(l=l) prof = Profiles() wfe = WavefrontError(n,m) print ('Your resolution is {}x{} with a pixel size of {}'.format(n,m,dpix)) A = prof.top_hat(38854/2, n, m, dpix=dpix, oversample=1) # Create an input grid, equivalent to the size of ELT-M1 x, y = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m)) # Create a ELT-M1 parabolic mirror M1 = ParabolicMirror(dx=0, dy=0, name='ELT-M1', CA=38854/2, f=679945) t_l = M1.phase_delay(x,y,l) U_in = A * t_l # Propagating the original, unperturbed, system and calculating encircled energy Exyz_og,xp,yp = fou.fresnel_propagation(U_in, x, y, dz=M1.f, show=False, outputCoordinates=oC) # Initiate psf fitter object to calculate the strehl ratios psf = PSFFitter(xp, yp, Exyz_og) StrehlList_zern_complete, StrehlList_psd_complete = [], [] rms_zern_list_complete, rms_psd_list_complete = [], [] for i in range(number_of_iterations): print('Iteration {}/{}'.format(i+1,number_of_iterations)) # Generate a PSD surface psd_surface = wfe.generate_height_map(midspat_rms, 2, 1) # Fit the PSD surface using zernike terms # Generating a normalized grid for the Zernike evaluation yz,xz = np.mgrid[-1:1:wfe.n*1j, -1:1:wfe.m*1j] _r = np.sqrt(xz**2 + yz**2) _mask = _r > 1 coeffic = wfe.fit_zernike_coefficients(psd_surface, modes=66) StrehlList_zern, StrehlList_psd = [], [] rms_zern_list, rms_psd_list = [], [] for rms in rmsList: # Make a new perturbed wavefront error with the required RMS coeffic = wfe.scale_zernike_coefficients(coeffic, rms) hzern = wfe.zernike_cartesian(coeffic, xz, yz) # Add the new wavefront error M1.add_wavefront_error(hzern) t_l = M1.phase_delay(x,y,l) U_in = A * t_l # Performing the propagation using Fresnel approximation to the focal plane. Exyz_zern,xp,yp = fou.fresnel_propagation(U_in, x, y, dz=M1.f, show=False, outputCoordinates=oC) Strehl_zern = np.abs(psf.get_coupling(Exyz_og, Exyz_zern))**2 set_trace() StrehlList_zern.append(Strehl_zern) #################################### # Add the new wavefront error factor = rms / midspat_rms h_psd = wfe.scale_rms(factor, data=psd_surface) # Add the new wavefront error M1.add_wavefront_error(h_psd) t_l = M1.phase_delay(x,y,l) U_in = A * t_l # Performing the propagation using Fresnel approximation to the focal plane. Exyz_psd,xp,yp = fou.fresnel_propagation(U_in, x, y, dz=M1.f, show=False, outputCoordinates=oC) Strehl_psd = np.abs(psf.get_coupling(Exyz_og, Exyz_psd))**2 StrehlList_psd.append(Strehl_psd) rms_zern = np.sqrt(np.mean(np.asarray(hzern)**2)) rms_psd = np.sqrt(np.mean(np.asarray(h_psd)**2)) print('RMS: {:.3f}, S_zern: {:.3f}, S_psd: {:.3f}, \ RMS_zern: {:.3f}, RMS_psd: {:.3f}'.format( rms*1e6, Strehl_zern, Strehl_psd, rms_zern*1e6, rms_psd*1e6)) rms_zern_list.append(rms_zern) rms_psd_list.append(rms_psd) StrehlList_zern_complete.append(StrehlList_zern) StrehlList_psd_complete.append(StrehlList_psd) rms_zern_list_complete.append(rms_zern_list) rms_psd_list_complete.append(rms_psd_list) mean_strehl_zern = np.mean(StrehlList_zern_complete, axis=0) err_strehl_zern = np.std(StrehlList_zern_complete, axis=0) mean_strehl_psd = np.mean(StrehlList_psd_complete, axis=0) err_strehl_psd = np.std(StrehlList_psd_complete, axis=0) rms_zern_list_complete = np.asarray(rms_zern_list_complete) rms_psd_list_complete = np.asarray(rms_psd_list_complete) # saveVariableToTXT(rmsList, 'rmsList.txt') # saveVariableToTXT(mean_strehl_zern, 'mean_strehl_zern.txt') # saveVariableToTXT(err_strehl_zern, 'err_strehl_zern.txt') # saveVariableToTXT(mean_strehl_psd, 'mean_strehl_psd.txt') # saveVariableToTXT(err_strehl_psd, 'err_strehl_psd.txt') # saveVariableToTXT(rms_psd_list_complete, 'rms_psd_list_complete.txt') # saveVariableToTXT(rms_zern_list_complete, 'rms_zern_list_complete.txt') # Normal plot fig = plt.figure() ax = fig.add_subplot(111) for i in range(number_of_iterations): # ax.plot(rmsList*1e6, StrehlList_zern_complete[i], 'y--', alpha=0.2) # ax.plot(rmsList*1e6, StrehlList_psd_complete[i], 'b--', alpha=0.2) ax.plot(rms_zern_list_complete[i]*1e6, StrehlList_zern_complete[i], 'y-', alpha=0.1) ax.plot(rms_psd_list_complete[i]*1e6, StrehlList_psd_complete[i], 'b-', alpha=0.1) ax.plot(rmsList*1e6, expected_strehl(rmsList,l), 'k-', label='Theory') # ax.errorbar(rmsList*1e6, mean_strehl_zern, yerr=err_strehl_zern, color='y', linestyle='-', label='Zernikes', capsize=3, markeredgewidth=1) # ax.errorbar(rmsList*1e6, mean_strehl_psd, yerr=err_strehl_psd, color='b', linestyle='-', label='PSD', capsize=3, markeredgewidth=1) ax.errorbar(rmsList*1e6, mean_strehl_zern, yerr=err_strehl_zern, color='y', linestyle='-', label='Zernikes', capsize=3, markeredgewidth=1) ax.errorbar(rmsList*1e6, mean_strehl_psd, yerr=err_strehl_psd, color='b', linestyle='-', label='PSD', capsize=3, markeredgewidth=1) ax.set_ylim([0,1]) ax.set_xlim([np.min(rmsList)*1e6, np.max(rmsList)*1e6]) ax.grid(True, alpha=0.3) ax.set_xlabel('RMS (nm)') ax.set_ylabel('Strehl') ax.legend() plt.tight_layout() # plt.savefig('strehl_verification.pdf', bbox_inches='tight', transparent=True) # Log plot slSquared = (rmsList / l)**2 err_strehl_zern = 0.434 * err_strehl_zern / mean_strehl_zern err_strehl_psd = 0.434 * err_strehl_psd / mean_strehl_psd fig = plt.figure() ax = fig.add_subplot(111) for i in range(number_of_iterations): ax.plot( (rms_zern_list_complete[i]/l)**2, np.log10(StrehlList_zern_complete[i]), 'y-', alpha=0.1) ax.plot( (rms_psd_list_complete[i]/l)**2, np.log10(StrehlList_psd_complete[i]), 'b-', alpha=0.1) ax.plot(slSquared, np.log10(expected_strehl(rmsList,l)), 'k-', label='Theory') ax.errorbar(slSquared, np.log10(mean_strehl_zern), yerr=err_strehl_zern, color='y', linestyle='-', label='Zernikes', capsize=3, markeredgewidth=1) ax.errorbar(slSquared, np.log10(mean_strehl_psd), yerr=err_strehl_psd, color='b', linestyle='-', label='PSD', capsize=3, markeredgewidth=1) # ax.set_ylim([0,1]) ax.set_xlim([np.min(slSquared), np.max(slSquared)]) ax.grid(True, alpha=0.3) ax.set_xlabel('$\\left(\\sigma/\\lambda\\right)^2$') ax.set_ylabel('$\\log_{10}(S)$') ax.legend() plt.tight_layout() # plt.savefig('strehl_verification_log.pdf', bbox_inches='tight', transparent=True) # Show the last psd and zernike surface that were generated psd_surface[_mask] = np.nan hzern[_mask] = np.nan fig2 = plt.figure(figsize=(10,4)) ax1 = fig2.add_subplot(121) ax2 = fig2.add_subplot(122) im1 = ax1.imshow(hzern, cmap='jet') im2 = ax2.imshow(psd_surface, cmap='jet') plt.tight_layout() # plt.savefig('wavefront_error_example_with_zern_fit.pdf', bbox_inches='tight', transparent=True) plt.show() # set_trace()