#!/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, '/home/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics') from fourierOptics import * from scipy import interpolate ''' This file simulates the paraxial optical model of the ELT-MICADO system with an ADC placed inside the collimated beam near the pupil. Surface form errors are added to the prisms to investigate the effect of WFEs on the PSF centroid. Author: Joost van den Born Contact: born@astron.nl ''' class ADC_surfaces: ''' Very specific class to generate and scale wavefront errors on the MICADO ADC prisms. ''' def __init__(self, x,y, clearAperture=125): self.n, self.m = x.shape self.clearAperture = clearAperture self.radius = self.clearAperture / 2 self.yz,self.xz = np.mgrid[x.min()/self.radius:x.max()/self.radius:self.n*1j, x.min()/self.radius:x.max()/self.radius:self.m*1j] self.zern_coefficients = None self.zern_surface = None self.psd_surface = None self.latest_surface = None def generate_ADC_wfe(self, zernRMS=20e-6, psdRMS=15e-6, zernModes=19): ''' We apply the wavefront error an each of the ADC prisms as a phase transformation, corresponding to a transmitted wavefront error. This wfe is separated into a zernike and psd component with a total rms of 25 nm zernike component - RMS: 20 nm psd component - RMS: 15 nm Total - RMS: 25 nm ''' print('Generating new ADC WFE...') wfe = WavefrontError(self.n,self.m) zern_rms = zernRMS # RMS of the zernike component of the WFE midspat_rms = psdRMS # RMS of the PSD component of the WFE self.psd_surface = wfe.generate_height_map(midspat_rms, 2, 1) zern_surface = wfe.generate_height_map(zern_rms, 2, 1) zern_coefficients = wfe.fit_zernike_coefficients(zern_surface, modes=zernModes) zern_coefficients[:3] = 0 # No defocus, tip or tilt self.zern_coefficients = wfe.scale_zernike_coefficients(zern_coefficients, zern_rms) # Scale the zernike modes to the clear aperture of the ADC self.zern_surface = wfe.zernike_cartesian(self.zern_coefficients, self.xz, self.yz, rescale=False) zern_surface = self.mask_surface(self.zern_surface) psd_surface = self.mask_surface(self.psd_surface) self.latest_surface = zern_surface + psd_surface return self.latest_surface def mask_surface(self, A): ''' Set the array values to 0 where the radius is > 1, without changing the original array. ''' # Rescale x and y _mask = np.sqrt(self.xz**2 + self.yz**2) > 1 B = A.copy() B[_mask] = 0 return B def scale_wfe(self, zernRMS=20e-5, psdRMS=15e-5): ''' Scaling of the original wfe, given the new zernike and psd RMS values that are desired. ''' print('Scaling ADC WFE...') _old_psd_rms = np.sqrt(np.mean(self.psd_surface**2)) _factor_psd = psdRMS / _old_psd_rms wfe = WavefrontError(self.n,self.m) new_psd_surface = wfe.scale_rms(_factor_psd, data=self.psd_surface) new_zern_coefficients = wfe.scale_zernike_coefficients(self.zern_coefficients, zernRMS) # Scale the zernike modes to the clear aperture of the ADC new_zern_surface = wfe.zernike_cartesian(new_zern_coefficients, self.xz, self.yz, rescale=False) new_zern_surface = self.mask_surface(new_zern_surface)#, _xz, _yz) new_psd_surface = self.mask_surface(new_psd_surface)#, _xz, _yz) self.latest_surface = new_zern_surface + new_psd_surface return self.latest_surface def fit_wfe_tilt(self,z=None): ''' Takes the WFE on the ADC, finds the clear aperture and fits only the tilt over that area. ''' if z is None: z = self.latest_surface _, _ix1 = find_nearest(self.xz,-1) _, _ix2 = find_nearest(self.xz, 1) _zselect = z[_ix1:_ix2,_ix1:_ix2] _wfe = WavefrontError(_zselect.shape[0], _zselect.shape[1]) _C = _wfe.fit_zernike_coefficients(_zselect, modes=3) # _C[0] = 0 _fit_surface = _wfe.zernike_cartesian(_C, self.xz, self.yz, rescale=False) _fit_surface = self.mask_surface(_fit_surface) return _fit_surface, _C if __name__ == '__main__': plt.rc('font', family='serif') plt.rc('axes', labelsize='large') iterations = 25 # Do the whole thing for 25 iterations for i in range(31,iterations+31): print('Iteration {:.0f}/{:.0f}'.format(i,iterations+31)) l = 1.35e-3 fou = FourierOptics(l=l) # MICADO paraxial system # Pick sampling parameters dpix = 50 n, m = 3001, 3001 print ('Your resolution is {}x{} with a pixel size of {}'.format(n,m,dpix)) # 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)) print('This means your grid is {} mm by {} mm, with each pixel {} mm.'.format(x.max()-x.min(), y.max()-y.min(), dpix)) scalingFactorList = [0.001, 0.01, 0.1, 0.5, 1, 2, 3, 5, 7, 10] ########################################################### #################### SETTING UP MICADO #################### ########################################################### # Define the ELT mask (the spiders and central obscuration) prof = Profiles() # We choose a top hat so that the fit in the pupil is easier ELT_MASK = prof.top_hat(19427, n,m,dpix)#.astype('float64') # ELT Mirror system combined into a single lens ELT_MIRROR = ParabolicMirror(dx=0, dy=0, name='ELT-paraxial', CA=38854, f=6.84022e5) ELT_FP = OpticalSurface(name='MCD-IFP') ELT_COL = ParaxialLens(CA=154.54, f=1418.94, name='MCD-COL') ELT_ADC = ADCObject() ELT_CAM = ParaxialLens(CA=89.2, f=4278.32701, name='MCD-CAM') ELT_DET = OpticalSurface(name='MCD-DET') MCD = OpticalTrain() MCD.add_object(ELT_MIRROR, dist=0) MCD.add_object(ELT_FP, dist=ELT_MIRROR.f) MCD.add_object(ELT_COL, dist=ELT_COL.f) ELT_ADC.add_to_optical_train(MCD, 1209.4718) MCD.add_object(ELT_CAM, dist=50) MCD.add_object(ELT_DET, dist=ELT_CAM.f) ########################################################### # Load propagation up to the ADC (so we don't need to compute this every time) Exyz3, x3, y3 = load_variable('./data/ADC_input_tophat_(1.35).pkl') x3 = x3.real y3 = y3.real E0, x0, y0 = load_variable('./data/MICADO_PSF_tophat_(1.35).pkl') # x3, y3 = load_variable('coordinates.pkl') ########################################################### #################### ADD WFE TO ADC ####################### ########################################################### # Do the analysis for different scaling factors popt_airy_list, popt_gauss_list, popt_moffat_list, popt_cust_list = [], [], [], [] pcov_airy_list, pcov_gauss_list, pcov_moffat_list, custom_var_list = [], [], [], [] popt_airy_list_tilt, popt_gauss_list_tilt, popt_moffat_list_tilt, popt_cust_list_tilt = [], [], [], [] pcov_airy_list_tilt, pcov_gauss_list_tilt, pcov_moffat_list_tilt, custom_var_list_tilt = [], [], [], [] # Add tilt results? strehl_list = [] strehl_list_tilt = [] # Generate the WFEs for scaling: wfe_p1 = ADC_surfaces(x3,y3) wfe_p2 = ADC_surfaces(x3,y3) wfe_p3 = ADC_surfaces(x3,y3) wfe_p4 = ADC_surfaces(x3,y3) wfe1 = wfe_p1.generate_ADC_wfe() wfe2 = wfe_p2.generate_ADC_wfe() wfe3 = wfe_p3.generate_ADC_wfe() wfe4 = wfe_p4.generate_ADC_wfe() for sc in scalingFactorList: for f in [1,2]: # 0 can be removed if we can load in E0 print('Scaling: {:.1f}, f: {:.1f}'.format(sc, f)) if f == 1: print('Adding WFE!') scaledZernRMS = 20e-6 * sc scaledPSDRMS = 15e-6 * sc ELT_ADC.p1.add_wavefront_error(wfe_p1.scale_wfe(zernRMS=scaledZernRMS, psdRMS=scaledPSDRMS)) ELT_ADC.p2.add_wavefront_error(wfe_p2.scale_wfe(zernRMS=scaledZernRMS, psdRMS=scaledPSDRMS)) ELT_ADC.p3.add_wavefront_error(wfe_p3.scale_wfe(zernRMS=scaledZernRMS, psdRMS=scaledPSDRMS)) ELT_ADC.p4.add_wavefront_error(wfe_p4.scale_wfe(zernRMS=scaledZernRMS, psdRMS=scaledPSDRMS)) if f == 2: print('Fitting WFE! (sc={:.2f}, f={:.2f})'.format(sc,f)) ELT_ADC.p1.add_wavefront_error(wfe1*0) ELT_ADC.p2.add_wavefront_error(wfe2*0) ELT_ADC.p3.add_wavefront_error(wfe3*0) ELT_ADC.p4.add_wavefront_error(wfe4*0) ########################################################### ########################################################### ###################### PROPAGATION ######################## ########################################################### coordinatesList = [ [-1.5, 1.5, -1.5, 1.5], [-200, 200, -200, 200], [-120, 120, -120, 120], [-120, 120, -120, 120], [-120, 120, -120, 120], [-120, 120, -120, 120], [-120, 120, -120, 120], [-2.5, 2.5, -2.5, 2.5] ] # if f==0: # xin, yin = x, y # print('Starting propagation...') # # From ELT-M1 to ELT-FP # # print('Propagating from ELT-M1 to ELT-IFP') # wf = ELT_MASK.astype('complex128') # wf *= ELT_MIRROR.phase_delay(x,y,l, paraxial=True) # Exyz, x, y = fou.fresnel_propagation(wf, xin, yin, dz=ELT_MIRROR.f, oversample=None, # show=False, save=False)#, outputCoordinates=coordinatesList[0], apply_interpolation=False) # # From ELT-FP to MCD-COL # # print('Propagating from ELT-IFP to MCD-COL') # Exyz2, x2, y2 = fou.fresnel_propagation(Exyz, x, y, dz=ELT_COL.f, oversample=None, # show=False, save=False, outputCoordinates=coordinatesList[1], apply_interpolation=False) # # From MCD-COL to MCD-ADC-P1 # # print('Propagating from MCD-COL to MCD-ADC-P1') # Exyz2 *= ELT_COL.phase_delay(x2,y2,l, paraxial=True) # Exyz3, x3, y3 = fou.plane_wave_propagation(Exyz2, x2, y2, dz=MCD.sequentialDistance[3], oversample=None, # show=False, save=False, outputCoordinates=None) # From MCD-ADC-P1 to MCD-ADC-P2 # print('Propagating from MCD-ADC-P1 to MCD-ADC-P2') Exyz3_ = Exyz3 * ELT_ADC.get_phase_delay(x3,y3,l, theta=None, paraxial=True, prism=1) Exyz4, x4, y4 = fou.plane_wave_propagation(Exyz3_, x3, y3, dz=MCD.sequentialDistance[4], oversample=None, show=False, save=False, outputCoordinates=None) # From MCD-ADC-P2 to MCD-ADC-P3 # print('Propagating from MCD-ADC-P2 to MCD-ADC-P3') Exyz4 *= ELT_ADC.get_phase_delay(x4,y4,l, theta=None, paraxial=True, prism=2) Exyz5, x5, y5 = fou.plane_wave_propagation(Exyz4, x4, y4, dz=MCD.sequentialDistance[5], oversample=None, show=False, save=False, outputCoordinates=None) # From MCD-ADC-P3 to MCD-ADC-P4 # print('Propagating from MCD-ADC-P3 to MCD-ADC-P4') Exyz5 *= ELT_ADC.get_phase_delay(x5,y5,l, theta=None, paraxial=True, prism=3) Exyz6, x6, y6 = fou.plane_wave_propagation(Exyz5, x5, y5, dz=MCD.sequentialDistance[6], oversample=None, show=False, save=False, outputCoordinates=None) # From MCD-ADC-P4 to MCD-CAM # print('Propagating from MCD-ADC-P4 to MCD-PUPIL') Exyz6 *= ELT_ADC.get_phase_delay(x6,y6,l, theta=None, paraxial=True, prism=4) Exyz7, x7, y7 = fou.plane_wave_propagation(Exyz6, x6, y6, dz=MCD.sequentialDistance[7], oversample=None, show=False, save=False, outputCoordinates=None) # MCD-CAM is also the pupil of the system: if f == 1: # Create a pupil plane object to fit the wavefront error pupilRadius = 40 # mm pupilPlane = ADC_surfaces(x7,y7, clearAperture=2*pupilRadius) # Convert the energy distribution into a distance 'lag' zWFE = np.angle(Exyz7)*l/(2*np.pi) # Remove any piston from the distance 'lag'. This does not add anything. pupilMask = np.sqrt(x7**2 + y7**2) < pupilRadius zWFE -= np.mean(zWFE[pupilMask]) pupilFittedWFE, pupilFittedC = pupilPlane.fit_wfe_tilt(z=zWFE) if f == 2: # Fitted pupil data into phase shift Exyz7 *= np.exp(1j*pupilFittedWFE * 2 * np.pi / l) # From MCD-CAM to MCD-DET # print('Propagating from MCD-CAM to MCD-DET') Exyz7 *= ELT_CAM.phase_delay(x7,y7,l, paraxial=True) Exyz8, x8, y8 = fou.fresnel_propagation(Exyz7, x7, y7, dz=ELT_CAM.f, oversample=None, show=False, save=False, outputCoordinates=None, apply_interpolation=False) print('Propagation finished!') ########################################################### ######################## FITTING ########################## ########################################################### print('Fitting the centroid...') fitter = PSFFitter(x8.real,y8.real,np.abs(Exyz8)**2) # if f==0: # E0 = Exyz8 if f == 1: airyFitted, popt_airy, pcov_airy = fitter.fit_airy_PSF(initialGuess=(0.03**2, 0.0, 0.0, 0.0, 0.03)) gaussianFitted, popt_gauss, pcov_gauss = fitter.fit_gaussian_PSF(initialGuess=(0.04, 0, 0, 0.01, 0.01, 0, 0)) # moffatFitted, popt_moffat, pcov_moffat = fitter.fit_moffat_PSF(initialGuess=(0.04, 0, 0, 0.01, 0.01, 0, 0, 4)) custFitted, popt_cust, cust_var, cust_coupling = fitter.fit_custom_field(E0, Exyz8, return_coupling=True) print('Strehl ratio: {:.4f}'.format(np.abs(cust_coupling)**2)) popt_airy_list.append(popt_airy) popt_gauss_list.append(popt_gauss) popt_cust_list.append(popt_cust) pcov_airy_list.append([ np.sqrt(pcov_airy[i][i]) for i in range(len(pcov_airy)) ]) pcov_gauss_list.append([ np.sqrt(pcov_gauss[i][i]) for i in range(len(pcov_gauss)) ]) custom_var_list.append(np.sqrt(cust_var)) strehl_list.append(np.abs(cust_coupling)**2) x0_cust, y0_cust = np.asarray(popt_cust).T r_cust = np.sqrt(x0_cust**2 + y0_cust**2) print (x0_cust, y0_cust) if f == 2: airyFitted_tilt, popt_airy_tilt, pcov_airy_tilt = fitter.fit_airy_PSF(initialGuess=(0.03**2, 0.0, 0.0, 0.0, 0.03)) gaussianFitted_tilt, popt_gauss_tilt, pcov_gauss_tilt = fitter.fit_gaussian_PSF(initialGuess=(0.04, 0, 0, 0.01, 0.01, 0, 0)) # moffatFitted, popt_moffat, pcov_moffat = fitter.fit_moffat_PSF(initialGuess=(0.04, 0, 0, 0.01, 0.01, 0, 0, 4)) if f==0: E0 = Exyz8 custFitted_tilt, popt_cust_tilt, cust_var_tilt, cust_coupling_tilt = fitter.fit_custom_field(E0, Exyz8, return_coupling=True) print('Strehl ratio (tilt fit): {:.4f}'.format(np.abs(cust_coupling_tilt)**2)) popt_airy_list_tilt.append(popt_airy_tilt) popt_gauss_list_tilt.append(popt_gauss_tilt) popt_cust_list_tilt.append(popt_cust_tilt) pcov_airy_list_tilt.append([ np.sqrt(pcov_airy_tilt[i][i]) for i in range(len(pcov_airy_tilt)) ]) pcov_gauss_list_tilt.append([ np.sqrt(pcov_gauss_tilt[i][i]) for i in range(len(pcov_gauss_tilt)) ]) custom_var_list_tilt.append(np.sqrt(cust_var_tilt)) strehl_list_tilt.append(np.abs(cust_coupling_tilt)**2) x0_cust_tilt, y0_cust_tilt = np.asarray(popt_cust_tilt).T r_cust_tilt = np.sqrt(x0_cust_tilt**2 + y0_cust_tilt**2) print (x0_cust_tilt, y0_cust_tilt) # Save the variables save_variable('./data/wfe scaling/other/airy_{:02d}.pkl'.format(i), [popt_airy_list, pcov_airy_list]) save_variable('./data/wfe scaling/other/gauss_{:02d}.pkl'.format(i), [popt_gauss_list, pcov_gauss_list]) save_variable('./data/wfe scaling/other/cust_{:02d}.pkl'.format(i), [popt_cust_list, custom_var_list, strehl_list]) save_variable('./data/wfe scaling/other/airy_tilt_{:02d}.pkl'.format(i), [popt_airy_list_tilt, pcov_airy_list_tilt]) save_variable('./data/wfe scaling/other/gauss_tilt_{:02d}.pkl'.format(i), [popt_gauss_list_tilt, pcov_gauss_list_tilt]) save_variable('./data/wfe scaling/other/cust_tilt_{:02d}.pkl'.format(i), [popt_cust_list_tilt, custom_var_list_tilt, strehl_list_tilt]) # popt_airy_list, pcov_airy_list = load_variable('./data/wfe scaling/airy.pkl') # popt_gauss_list, pcov_gauss_list = load_variable('./data/wfe scaling/gauss.pkl') # popt_cust_list, custom_var_list, strehl_list = load_variable('./data/wfe scaling/cust.pkl') # amplitude_airy, x0_airy, y0_airy, R_airy, offset_airy = np.asarray(popt_airy_list).T # A_gauss, x0_gauss, y0_gauss, sigma_x_gauss, sigma_y_gauss, theta_gauss, offset_gauss = np.asarray(popt_gauss_list).T # # A_moffat, x0_moffat, y0_moffat, sigma_x_moffat, sigma_y_moffat, theta_moffat, offset_moffat, beta_moffat = np.asarray(popt_moffat_list).T # x0_cust, y0_cust = np.asarray(popt_cust_list).T # Check if this is correct! # # Errors in the optimum fit parameters for all scaled surface errors # amplitude_airy_err, x0_airy_err, y0_airy_err, R_airy_err, offset_airy_err = np.asarray(pcov_airy_list).T # A_gauss_err, x0_gauss_err, y0_gauss_err, sigma_x_gauss_err, sigma_y_gauss_err, theta_gauss_err, offset_gauss_err = np.asarray(pcov_gauss_list).T # # A_moffat_err, x0_moffat_err, y0_moffat_err, sigma_x_moffat_err, sigma_y_moffat_err, theta_moffat_err, offset_moffat_err, beta_moffat_err = np.asarray(pcov_moffat_list).T # x0_cust_err, y0_cust_err = np.asarray(custom_var_list).T # # Distance from original centroid at (0,0) # r_airy = np.sqrt(x0_airy**2+y0_airy**2) # r_err_airy = np.sqrt( (x0_airy**2*x0_airy_err**2+y0_airy**2*y0_airy_err**2) / (x0_airy**2+y0_airy**2)) # r_gauss = np.sqrt(x0_gauss**2+y0_gauss**2) # r_err_gauss = np.sqrt( (x0_gauss**2*x0_gauss_err**2+y0_gauss**2*y0_gauss_err**2) / (x0_gauss**2+y0_gauss**2)) # # r_moffat = np.sqrt(x0_moffat**2+y0_moffat**2) # # r_err_moffat = np.sqrt( (x0_moffat**2*x0_moffat_err**2+y0_moffat**2*y0_moffat_err**2) / (x0_moffat**2+y0_moffat**2)) # r_cust = np.sqrt(x0_cust**2 + y0_cust**2) # r_err_cust = np.sqrt( (x0_cust**2*x0_cust_err**2+y0_cust**2*y0_cust_err**2) / (x0_cust**2+y0_cust**2)) ########################################################### ######################## PLOTTING ######################### ########################################################### # Plot the PSF shape # _xip = x8.real # _yip = y8.real # Exyz = Exyz8 # save = 'ELT-MCD-PSF_at_{}_WFE({}nm).pdf'.format((l*1000),(f*25)) # _fig = plt.figure(figsize=(16,6)) # _ax = _fig.add_subplot(121) # _im = _ax.imshow(np.abs(Exyz.T)**2, cmap='hot', # extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()], norm=LogNorm(vmin=1e-8)) # _ax.set_xlim([-10, 10]) # _ax.set_ylim([-10, 10]) # _ax.set_xlabel('x [mm]') # _ax.set_ylabel('y [mm]') # _fig.colorbar(_im, ax=_ax, label='Intensity') # _ax2 = _fig.add_subplot(122) # _im2 = _ax2.imshow(np.angle(np.angle(Exyz.T)), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()]) # _ax2.set_xlim([-10, 10]) # _ax2.set_ylim([-10, 10]) # _ax2.set_xlabel('x [mm]') # _ax2.set_ylabel('y [mm]') # _fig.colorbar(_im2, ax=_ax2, label='Phase') # plt.tight_layout() # plt.savefig(save, bbox_inches='tight', transparent=True) # plt.show() # Plot the centroid position RMSlist = [25*f for f in fList] fig = plt.figure() ax1 = fig.add_subplot(121) ax2 = fig.add_subplot(122) ax1.errorbar(RMSlist, r_gauss, yerr=r_err_gauss, color='green', marker='o', linestyle='--', label='Gaussian', capsize=3, markeredgewidth=1) ax1.errorbar(RMSlist, r_airy, yerr=r_err_airy, color='blue', marker='o', linestyle='--', label='Airy', capsize=3, markeredgewidth=1) # ax1.errorbar(RMSlist, r_moffat, yerr=r_err_moffat, color='green', marker='o', linestyle='--', label='Moffat', capsize=3, markeredgewidth=1) ax1.errorbar(RMSlist, r_cust, yerr=r_err_cust, color='orange', marker='o', linestyle='--', label='Custom', capsize=3, markeredgewidth=1) # ax1.axhline(1.875e-4, color='red', linestyle='-', label='50 $\mu$as @ 4 mas/pix') ax1.set_xscale('log') ax1.set_yscale('log') ax1.set_xlabel('RMS WFE (nm)') ax1.set_ylabel('Distance') ax1.legend() ax1.grid(True, alpha=0.3, which='major') ax1.grid(True, alpha=0.1, which='minor') ax2.errorbar(x0_airy, y0_airy, xerr=x0_airy_err, yerr=y0_airy_err, color='k', marker='o', linestyle='--', label='Airy') ax2.errorbar(x0_cust, y0_cust, xerr=x0_cust_err, yerr=y0_cust_err, color='r', marker='o', linestyle='--', label='Custom') ax2.set_aspect('equal') ax2.grid(True, alpha=0.3) plt.tight_layout() plt.show() ###########################################################