#!/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 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)#, self.xz, self.yz) psd_surface = self.mask_surface(self.psd_surface)#, self.xz, self.yz) 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 if __name__ == '__main__': plt.rc('font', family='serif') plt.rc('axes', labelsize='large') 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)) fList = [1]#[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() # ELT_MASK = prof.top_hat(19427, n,m,dpix)#.astype('float64') # ELT_MASK = prof.obscured_top_hat(4000, 500, n, m, dpix) # ELT_MASK = prof.ELT_pupil(n,m,dpix) # 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) # Print lens data for verification MCD.print_lens_data() ########################################################### # Load propagation up to the ADC (so we don't need to compute this every time) Exyz3, x3, y3 = load_variable('./data/ADC_input_(1.35).pkl') x3 = x3.real y3 = y3.real E0, x0, y0 = load_variable('./data/MICADO_PSF_(1.35).pkl') ########################################################### #################### ADD WFE TO ADC ####################### ########################################################### # adc_wfe = ADC_surfaces(x3,y3,125) # wfe_p1 = adc_wfe.generate_ADC_wfe() # newsurf = adc_wfe.scale_wfe() wfe_p1 = ADC_surfaces(x3,y3) wfe_p2 = ADC_surfaces(x3,y3) wfe_p3 = ADC_surfaces(x3,y3) wfe_p4 = ADC_surfaces(x3,y3) # wfe_p2 = generate_ADC_wfe(x3, y3) # wfe_p3 = generate_ADC_wfe(x3, y3) # wfe_p4 = generate_ADC_wfe(x3, y3) wfe_p1.generate_ADC_wfe() set_trace() wfe_p2.generate_ADC_wfe() wfe_p3.generate_ADC_wfe() wfe_p4.generate_ADC_wfe() # 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 = [], [], [], [] strehl_list = [] for f in fList: print('Factor set to {:.1f}'.format(f)) scaledZernRMS = 20e-6 * f scaledPSDRMS = 15e-6 * f 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)) ########################################################### ########################################################### ###################### 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] ] print('Starting propagation...') # 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-CAM') 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) # 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) # # The result here shifts the PSF to the side, due to the phase tilts of the ADC # # surfaces. We find the locations of the maximum and redo the the fresnel propagation # # around these coordinates. # print('Zooming in on the PSF location...') # imax_1, imax_2 = find_maximum(np.abs(Exyz8)) # x8max = x8[imax_1, imax_2] # y8max = y8[imax_1, imax_2] # newCoordinates = np.asarray(coordinatesList[7]) + np.array([x8max, x8max, y8max, y8max]) # newCoordinates = newCoordinates.real.tolist() # Exyz9, x9, y9 = fou.fresnel_propagation(Exyz7, x7, y7, dz=ELT_CAM.f, oversample=None, # show=False, save=False, outputCoordinates=newCoordinates, apply_interpolation=True) print('Propagation finished!') ########################################################### ######################## FITTING ########################## ########################################################### print('Fitting the centroid...') fitter = PSFFitter(x8.real,y8.real,np.abs(Exyz8)**2) 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) # Save the variables save_variable('./data/wfe scaling/airy_5.pkl', [popt_airy_list, pcov_airy_list]) save_variable('./data/wfe scaling/gauss_5.pkl', [popt_gauss_list, pcov_gauss_list]) save_variable('./data/wfe scaling/cust_5.pkl', [popt_cust_list, custom_var_list, strehl_list]) # 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() ###########################################################