diff --git a/scripts/ELT_off_axis_test.py b/scripts/ELT_off_axis_test.py new file mode 100644 index 0000000..f45854d --- /dev/null +++ b/scripts/ELT_off_axis_test.py @@ -0,0 +1,88 @@ +#!/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 pickle +import sys +sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics') +from fourierOptics import * + + +''' + @ Author: Joost van den Born +''' +def saveVariable(var, filename): + # Saving the objects: + with open(filename, 'wb') as f: # Python 3: open(..., 'wb') + pickle.dump(var, f) + +def loadVariable(filename): + # Getting back the objects: + with open(filename, 'rb') as f: # Python 3: open(..., 'rb') + obj = pickle.load(f) + return obj + +def phaseTilt(x,y,tx,ty): + n,m = x.shape + y,x = np.mgrid[-1:1:n*1j,-1:1:m*1j] + phaseTilt = x * np.tan(np.deg2rad(tx)) + y * np.tan(np.deg2rad(ty)) + return phaseTilt + + +if __name__ == '__main__': + plt.rc('font', family='serif') + plt.rc('axes', labelsize='large') + + lList = np.linspace(1.49e-3, 1.78e-3, 29) + Exyz_list = [] + # Pick a wavelength + l = 1.35e-3 + + # Pick sampling parameters + dpix = 160 + n, m = 1001, 1001 + # oC = [-.75, .75, -.75, .75] + # oC = [-1.65, -0.15,-1.65, -0.15] + oC = [-3,3,-3,3] + # Initiate the profiles class and pick a profile (tophat, obscured tophat or spider) + prof = Profiles() + A = prof.ELT_pupil(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)) + # M1 = ParaxialLens(f=6.84022e5, CA=38854, name='ELT-paraxial') + M1 = ParabolicMirror(dx=0, dy=0, name='ELT-paraxial', CA=38854, f=6.84022e5) + M1.add_wavefront_error(phaseTilt(x,y,tx=7.5e-3, ty=7.5e-3)) + + for l in lList: + print(l) + # Initiate Fourier Optics + fou = FourierOptics(l=l) + + + print ('Your resolution is {}x{} with a pixel size of {}'.format(n,m,dpix)) + + + # Create a lens with a focal length ~100 mm, with a clear aperture radius of 4 mm and thickness of 1 mm. + 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) + Exyz_list.append(Exyz_og) + + Exyz = np.asarray(Exyz_list).sum(axis=0) + plt.imshow(np.abs(Exyz), cmap='hot');plt.colorbar();plt.show() + + + + + + + + + + diff --git a/scripts/ELT_wfe_analysis.py b/scripts/ELT_wfe_analysis.py new file mode 100644 index 0000000..7742ef7 --- /dev/null +++ b/scripts/ELT_wfe_analysis.py @@ -0,0 +1,279 @@ +#!/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 pickle +import sys +sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics') +from fourierOptics import * + + +''' + @ Author: Joost van den Born +''' +def saveVariable(var, filename): + # Saving the objects: + with open(filename, 'wb') as f: # Python 3: open(..., 'wb') + pickle.dump(var, f) + +def loadVariable(filename): + # Getting back the objects: + with open(filename, 'rb') as f: # Python 3: open(..., 'rb') + obj = pickle.load(f) + return obj + +def phaseTilt(x,y,tx,ty): + phaseTilt = x * np.tan(np.deg2rad(tx)) + y * np.tan(np.deg2rad(ty)) + return phaseTilt + + +if __name__ == '__main__': + plt.rc('font', family='serif') + plt.rc('axes', labelsize='large') + + # Pick a wavelength + l = 1.35e-3 + + # Initiate Fourier Optics + fou = FourierOptics(l=l) + + # Pick sampling parameters + dpix = 160 + n, m = 1001, 1001 + oC = [-.75, .75, -.75, .75] + 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() + # Aideal = prof.top_hat(38854/2, n, m, dpix=dpix, oversample=1) + A = prof.ELT_pupil(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)) + + # Add PSD wavefront error to the lens + wfe = WavefrontError(n,m) + wfe.generate_height_map(5e-5, 2, 1) # 50 nm RMS for standard + h_data = wfe.data + + + # Setting up empty lists for fitting process + popt_airy_list, popt_cust_list = [], [] + pcov_airy_list, custom_var_list = [], [] + strehl_list = [] + Exyzlist, PSDlist, surfacelist, EEFlist = [], [], [], [] + + + # Create a lens with a focal length ~100 mm, with a clear aperture radius of 4 mm and thickness of 1 mm. + M1 = ParabolicMirror(dx=0, dy=0, name='ELT-paraxial', CA=38854, f=6.84022e5) + # M1.add_wavefront_error(phaseTilt(x,y,tx=7.5, ty=7.5)) + 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) + eef_og, r_og = fou.encircled_energy_function(np.abs(Exyz_og)**2, xp, yp) + + + plt.imshow(np.abs(Exyz_og), cmap='hot');plt.colorbar();plt.show() + set_trace() + zernikeAnalysis = True + PSDAnalysis = True + + if zernikeAnalysis: + # For some reason we need to reset the WFE object. + wfe = WavefrontError(n,m) + wfe.data = h_data + psd1d_og = wfe.calculate_1D_PSD(data=h_data) + + # For each Zernike mode, perform analysis. + # Don't touch the RMS values, it will converge to some value + # If it goes as expected you'll see some form of convergence, + # possibly a spiral, of the centroid location as you add more + # and more modes. + + # Calculate the maximum mode that we want to fit the zernike modes to. + nmax, mmax = 80, 80 + zernikeLimit = int((nmax*(nmax+2)+mmax)/2) + + # Fit the zernike coefficients for all modes up to zernikeLimit + print('Fitting Zernike up to j={:.0f}'.format(zernikeLimit)) + zernikeCoefficients = wfe.fit_zernike_coefficients(h_data, modes=zernikeLimit) + + + # Result lists: + x0zern, y0zern = [], [] + x0zern_var, y0zern_var = [], [] + zernCouplingList = [] + eef_zern_list = [] + zernImageResponse = [] + psd1d_list_zern = [] + + for nmax_i in range(nmax): + print (' Zernike analysis: {0}/{1}'.format(nmax_i, nmax), end="\r") + zernikeModeMax = int((nmax_i*(nmax_i+2)+nmax_i)/2+1) + + # The zernike coefficients for this particular mode + zC = zernikeCoefficients[:zernikeModeMax] + xz, yz = np.mgrid[-1:1:n*1j, -1:1:m*1j] + wfeFilteredZern = wfe.zernike_cartesian(zC, xz, yz) + + # Adding the wavefront error + M1.add_wavefront_error(wfeFilteredZern) + wfe_t_l = M1.phase_delay(x,y,l) + U_in = A * wfe_t_l + + # Performing the propagation using Fresnel approximation to the focal plane. + Exyz_zern,xz,yz = fou.fresnel_propagation(U_in, x, y, dz=M1.f, show=False, outputCoordinates=oC) + eef_zern, _ = fou.encircled_energy_function(np.abs(Exyz_zern)**2, xz, yz, originalPSF=np.abs(Exyz_og)**2) + + # Fitting the centroid location using the coupling integral + # Set up fitting object + psfFit_zern = PSFFitter(xz.real,yz.real,np.abs(Exyz_zern)**2) + custFitted, popt_cust, cust_var, cust_coupling = psfFit_zern.fit_custom_field(Exyz_og, Exyz_zern, return_coupling=True) + + # Returning the found PSD curve + psd1d_list_zern.append(wfe.calculate_1D_PSD(data=wfeFilteredZern)) + + x0zern.append(popt_cust[0]) + y0zern.append(popt_cust[1]) + x0zern_var.append(cust_var[0]) + y0zern_var.append(cust_var[1]) + zernCouplingList.append(cust_coupling) + eef_zern_list.append(eef_zern.tolist()) + zernImageResponse.append(Exyz_zern) + + saveVariable(x0zern, './data/x0zern.pck') + saveVariable(y0zern, './data/y0zern.pck') + saveVariable(x0zern_var, './data/x0zern_var.pck') + saveVariable(y0zern_var, './data/y0zern_var.pck') + saveVariable(zernCouplingList, './data/zernCouplingList.pck') + saveVariable(eef_zern_list, './data/eef_zern_list.pck') + saveVariable(zernImageResponse, './data/zernImageResponse.pck') + saveVariable(psd1d_list_zern, './data/psd1d_list_zern.pck') + saveVariable(zernikeCoefficients, './data/zernikeCoefficients.pcks') + + if PSDAnalysis: + # For each spatial frequency, perform the analysis + # Don't touch the RMS value. It will converge to some value + + # For some reason we need to reset the WFE object. + wfe = None + wfe = WavefrontError(n,m) + wfe.data = h_data + psd1d_og = wfe.calculate_1D_PSD(data=h_data) + + # Result lists: + x0psd, y0psd = [], [] + x0psd_var, y0psd_var =[], [] + psdCouplingList = [] + eef_psd_list = [] + psdImageResponse = [] + frequencyRange = np.asarray([1,2,3,4,5,6,7,8,9] + np.linspace(10,int(np.sqrt(2)*n/2+1),70).tolist()) + psd1d_list_psd = [] + + for lowpassFrequency in frequencyRange: + i1 = frequencyRange.tolist().index(lowpassFrequency) + i2 = len(frequencyRange) + print (' Power Spectral Density analysis: {0}/{1}'.format(i1, i2), end="\r") + # Filter the spatial frequencies + wfeFilteredPSD = wfe.set_lowpass_filter(lowpassFrequency) + + # Adding the wavefront error + M1.add_wavefront_error(wfeFilteredPSD) + wfe_t_l = M1.phase_delay(x,y,l) + U_in = A * wfe_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) + eef_psd, _ = fou.encircled_energy_function(np.abs(Exyz_psd)**2, xp, yp, originalPSF=np.abs(Exyz_og)**2) + + # Fitting the centroid location using the coupling integral + # Set up fitting object + psfFit_psd = PSFFitter(xp.real,yp.real,np.abs(Exyz_psd)**2) + custFitted, popt_cust, cust_var, cust_coupling = psfFit_psd.fit_custom_field(Exyz_og, Exyz_psd, return_coupling=True) + + # Returning the found PSD curve + psd1d_list_psd.append(wfe.calculate_1D_PSD(data=wfeFilteredPSD)) + + x0psd.append(popt_cust[0]) + y0psd.append(popt_cust[1]) + x0psd_var.append(cust_var[0]) + y0psd_var.append(cust_var[1]) + psdCouplingList.append(cust_coupling) + eef_psd_list.append(eef_psd) + psdImageResponse.append(Exyz_psd) + + saveVariable(x0psd, './data/x0psd.pck') + saveVariable(y0psd, './data/y0psd.pck') + saveVariable(x0psd_var, './data/x0psd_var.pck') + saveVariable(y0psd_var, './data/y0psd_var.pck') + saveVariable(psdCouplingList, './data/psdCouplingList.pck') + saveVariable(eef_psd_list, './data/eef_psd_list.pck') + saveVariable(psdImageResponse, './data/psdImageResponse.pck') + saveVariable(psd1d_list_psd, './data/psd1d_list_psd.pck') + + + + # Plot some of the results: + if PSDAnalysis: + # Power Spectral Density distribution for PSD analysis + fig = plt.figure() + ax = fig.add_subplot(111) + ax.loglog(psd1d_og, 'r-') # The original PSD curve + # The PSD curve for each lowpass filtered PSD WFE + for i in range(len(frequencyRange)): + ax.loglog(psd1d_list_psd[i], 'b-', label='data') + ax.grid(True, alpha=0.3) + ax.set_xlabel('Spatial frequency') + ax.set_ylabel('PSD') + plt.show(block=False) + + if zernikeAnalysis: + # Power Spectral Density distribution for Zernike Analysis + fig = plt.figure() + ax = fig.add_subplot(111) + # The original PSD curve + ax.loglog(psd1d_og, 'r-') + # The PSD curve for each of the zernike fits + for i in range(nmax): + ax.loglog(psd1d_list_zern[i], 'b-', label='data') + ax.grid(True, alpha=0.3) + ax.set_xlabel('Spatial frequency') + ax.set_ylabel('PSD') + plt.show() + + + + # Plot of the centroid shift + fig = plt.figure() + ax = fig.add_subplot(111) + + if zernikeAnalysis: + ax.plot(x0zern, y0zern, 'b-', label='data') + ax.plot(x0zern, y0zern, 'bx', label='data') + ax.plot(x0zern[0], y0zern[0], 'rx') + + if PSDAnalysis: + ax.plot(x0psd, y0psd, 'r-', label='data') + ax.plot(x0psd, y0psd, 'rx', label='data') + ax.plot(x0psd[0], y0psd[0], 'bx') + + ax.grid(True, alpha=0.3) + ax.set_xlabel('x [mm]') + ax.set_ylabel('y [mm]') + plt.show() + + + + + + + + + + diff --git a/scripts/MICADO_ADC_with_WFE.py b/scripts/MICADO_ADC_with_WFE.py new file mode 100644 index 0000000..6e855ad --- /dev/null +++ b/scripts/MICADO_ADC_with_WFE.py @@ -0,0 +1,392 @@ +#!/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() + + ########################################################### diff --git a/scripts/MICADO_ADC_with_WFE_iterations.py b/scripts/MICADO_ADC_with_WFE_iterations.py new file mode 100644 index 0000000..59e5aee --- /dev/null +++ b/scripts/MICADO_ADC_with_WFE_iterations.py @@ -0,0 +1,399 @@ +#!/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, 'D:\\Users\\Joost\\Desktop\\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') + + data_airy, data_gauss, data_cust = [],[],[] + + iterations = 25 + for i in range(iterations): + 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 = [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() + 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 + data_airy.append([popt_airy_list, pcov_airy_list]) + data_gauss.append([popt_gauss_list, pcov_gauss_list]) + data_cust.append([popt_cust_list, custom_var_list, strehl_list]) + + save_variable('.\\data\\wfe scaling\\airy_complete.pkl', data_airy) + save_variable('.\\data\\wfe scaling\\gauss_complete.pkl', data_gauss) + save_variable('.\\data\\wfe scaling\\cust_complete.pkl', data_cust) + + # 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') + + set_trace() + + 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() + + ########################################################### diff --git a/scripts/MICADO_paraxial_system.py b/scripts/MICADO_paraxial_system.py new file mode 100644 index 0000000..87363ed --- /dev/null +++ b/scripts/MICADO_paraxial_system.py @@ -0,0 +1,179 @@ +#!/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. + + Author: Joost van den Born + Contact: born@astron.nl +''' + + + +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)) + + + ########################################################### + #################### 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() + + ########################################################### + + + + + + ########################################################### + ###################### 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] + ] + + xin, yin = x, y + zList = np.linspace(ELT_MIRROR.f-100, ELT_MIRROR.f+100, 15).tolist() + 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 *= 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) + + # save_variable('./data/MICADO_PSF_(1.35).pkl', [Exyz8, x8, y8]) + + # 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() + + save = 'MICADO_PSF.pdf' + Exyz9, x9, y9 = fou.fresnel_propagation(Exyz7, x7, y7, dz=ELT_CAM.f, oversample=None, + show=True, save=save, outputCoordinates=newCoordinates, apply_interpolation=True) + print('Propagation finished!') + # save_variable('./data/MICADO_PSF_(1.35)_interpolated.pkl', [Exyz9, x9, y9]) + ########################################################### + + + + + + + + + diff --git a/scripts/MICADO_paraxial_system_dispersion.py b/scripts/MICADO_paraxial_system_dispersion.py new file mode 100644 index 0000000..91595b9 --- /dev/null +++ b/scripts/MICADO_paraxial_system_dispersion.py @@ -0,0 +1,253 @@ +#!/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 + +def perform_padding(A, xs, ys): + assert A.shape[0] <= xs and A.shape[1] <= ys, 'New size is smaller than array size.' + x_margin = int((xs - A.shape[0])/2) + y_margin = int((ys - A.shape[1])/2) + return np.pad(A, (x_margin, y_margin)) + +def find_nearest(array, value): + array = np.asarray(array) + diff = (np.abs(array - value)).argmin() + idx = np.unravel_index(diff, array.shape) + return idx + +def find_maximum(array): + idx = np.unravel_index(array.argmax(), array.shape) + return idx + +def select_subsection(x, y , outputCoordinates): + select_1 = find_nearest(x, outputCoordinates[0])[1] + select_2 = find_nearest(x, outputCoordinates[1])[1] + select_3 = find_nearest(y, outputCoordinates[2])[0] + select_4 = find_nearest(y, outputCoordinates[3])[0] + return [select_1, select_2, select_3, select_4] + +def perform_extrapolation(x, y, n, m): + # Determine properties of x,y meshgrid + xshape = x.shape[0] + yshape = y.shape[0] + dpix = (np.abs(np.min(x)) + np.abs(np.max(x))) / (xshape-1) + + # Extrapolate from, and include, the old x,y meshgrid + xnew = np.linspace(-(n-1)/2*dpix, (n-1)/2*dpix, n) + ynew = np.linspace(-(m-1)/2*dpix, (m-1)/2*dpix, m) + xnew, ynew = np.meshgrid(xnew, ynew) + return xnew, ynew + +if __name__ == '__main__': + plt.rc('font', family='serif') + plt.rc('axes', labelsize='large') + + + # Dispersion levels: + lList = np.linspace(1.49e-3, 1.78e-3, 2).tolist() + dispMax = 17.9e-3 / 3600 # dispersion in degrees + dispersion = np.linspace(0, dispMax, len(lList)).tolist() + Exyz_list = [] + + # Input parameters: + dpix = 50 + n, m = 3001, 3001 + x, y = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m)) + print ('Your resolution is {}x{} with a pixel size of {}'.format(n,m,dpix)) + print('This means your grid is {} mm by {} mm, with each pixel {} mm.'.format(x.max()-x.min(), y.max()-y.min(), dpix)) + + ########################################################### + #################### 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) + + # plt.imshow(ELT_MASK);plt.colorbar();plt.show() + + + + + + + + ########################################################### + + for l in lList: + i = lList.index(l) + print('Wavelength is {:.6f} mm. ({:02d}/{}).'.format(l,lList.index(l)+1, len(lList))) + x, y = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m)) + + fou = FourierOptics(l=l) + wfe = WavefrontError(n,m) + dispersion_tilt = wfe.generate_tilt(x, y, tx=0, ty=dispersion[i]) + + # ELT Mirror system combined into a single lens + ELT_MIRROR = ParabolicMirror(dx=0, dy=0, name='ELT-paraxial', CA=38542, f=6.84022e5) + + ELT_FP = OpticalSurface(name='MCD-IFP') + + ELT_COL = ParaxialLens(CA=154.54, f=1418.94, name='MCD-COL') + + ELT_ADC = ADCObject() + + # optAngle = ELT_ADC.find_optimum_angle(17.9, lList[0], lList[-1]) + # print(optAngle) + print('ADC angle now set to {:.4f}'.format(ELT_ADC.theta)) + # ELT_ADC.set_orientation(180) + + 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() + + ########################################################### + ###################### 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], + [-1.5, 1.5, -1.5, 1.5] + ] + + xin, yin = x, y + + print('Starting propagation...') + + # From ELT-M1 to ELT-FP + print('Propagating from ELT-M1 to ELT-IFP', end='\r') + wf = ELT_MASK.astype('complex128') + ELT_MIRROR.add_wavefront_error(dispersion_tilt) + 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) + save_variable('./data/ELT_FP_{}.pkl'.format(l*1000), [Exyz, x, y]) + set_trace() + # From ELT-FP to MCD-COL + print('Propagating from ELT-IFP to MCD-COL', end='\r') + 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', end='\r') + 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', end='\r') + 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', end='\r') + 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', end='\r') + 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', end='\r') + 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', end='\r') + 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...', end='\r') + # 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!') + + + Exyz_list.append(np.abs(Exyz8)**2) + fou = None + del fou + wfe = None + del wfe + MCD = None + del MCD + wf = None + del wf + ELT_MIRROR = None + del ELT_MIRROR + ELT_COL = None + del ELT_COL + ELT_ADC = None + del ELT_ADC + ELT_CAM = None + del ELT_CAM + ELT_DET = None + del ELT_DET + ########################################################### + + + Exyz_total = np.asarray(Exyz_list).sum(axis=0) + + _xip = x8 + _yip = y8 + # save = 'ELT-MCD-Hband_dispersion_corrected.pdf' + _fig = plt.figure() + _ax = _fig.add_subplot(111) + _im = _ax.imshow(Exyz_total.T, cmap='hot', + extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()], norm=LogNorm(vmin=2e-7)) + _ax.set_xlim([_xip.min(), _xip.max()]) + _ax.set_ylim([_yip.min(), _yip.max()]) + _ax.set_xlabel('x [mm]') + _ax.set_ylabel('y [mm]') + _fig.colorbar(_im, ax=_ax, label='Intensity') + plt.tight_layout() + # plt.savefig(save, bbox_inches='tight', transparent=True) + plt.show() + + set_trace() + + + + + diff --git a/scripts/MICADO_tilt_fit.py b/scripts/MICADO_tilt_fit.py new file mode 100644 index 0000000..9372eb7 --- /dev/null +++ b/scripts/MICADO_tilt_fit.py @@ -0,0 +1,408 @@ +#!/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) + 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) + _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') + + 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() + wfe_p2.generate_ADC_wfe() + wfe_p3.generate_ADC_wfe() + wfe_p4.generate_ADC_wfe() + fit1, _ = wfe_p1.fit_wfe_tilt() + fit2, _ = wfe_p2.fit_wfe_tilt() + fit3, _ = wfe_p3.fit_wfe_tilt() + fit4, _ = wfe_p4.fit_wfe_tilt() + + # 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 [0,1]: + print('Factor set to {:.1f}'.format(f)) + ELT_ADC.p1.add_wavefront_error([wfe_p1.latest_surface, fit1][f]) + ELT_ADC.p2.add_wavefront_error([wfe_p2.latest_surface, fit2][f]) + ELT_ADC.p3.add_wavefront_error([wfe_p3.latest_surface, fit3][f]) + ELT_ADC.p4.add_wavefront_error([wfe_p4.latest_surface, fit4][f]) + ########################################################### + + + + + + ########################################################### + ###################### 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) + + print (x0_cust, y0_cust) + + # 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() + + ########################################################### diff --git a/scripts/MICADO_tilt_fit_in_pupil.py b/scripts/MICADO_tilt_fit_in_pupil.py new file mode 100644 index 0000000..7b34eeb --- /dev/null +++ b/scripts/MICADO_tilt_fit_in_pupil.py @@ -0,0 +1,474 @@ +#!/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() + + ########################################################### diff --git a/scripts/annular_mask_lens.py b/scripts/annular_mask_lens.py new file mode 100644 index 0000000..4e38f11 --- /dev/null +++ b/scripts/annular_mask_lens.py @@ -0,0 +1,135 @@ +#!/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 +from fourierOptics 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.obscured_top_hat(2.5, 2.3, n, m, dpix=dpix, oversample=1) + plt.imshow(A); plt.show() + + + # 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)) + + + + # 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', 50, 1e99, 1, 4, name='L1') + print('Focal length of lens {} is {:.1f} mm.'.format(L1.name, L1.f)) + + # Calculate the phase delay introduced by the lens + t_l = L1.phase_delay(x,y,l, paraxial=True) + 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) + + + + E1, fx1, fy1 = fou.modulation_transfer_function(np.abs(Exyz_og)**2, x,y,dpix) + plt.figure() + plt.plot(fx1[1000:], E1[1000, 1000:],'b-') + plt.xlim(0,1.67) + plt.ylim(0,1) + plt.xlabel('Spatial frequency (cycles/aperture)') + plt.ylabel('MTF') + plt.grid() + plt.show() + + + + + # The following produces a cross-section of the propagation for various z, so that you can see how the beam evolves. + if True: + Ics = [] + Ics_x = [] + Isum = [] + xmax = [] + zList = np.linspace(10,150,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(111) + 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))) + + + diff --git a/scripts/artificial_speckles_test.py b/scripts/artificial_speckles_test.py new file mode 100644 index 0000000..8e0810e --- /dev/null +++ b/scripts/artificial_speckles_test.py @@ -0,0 +1,64 @@ +#!/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 sys +sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics') +from fourierOptics import * + +plt.rc('font', family='serif') +plt.rc('axes', labelsize='large') + +# Parameters +dpix = 160 +n, m = 1001, 1001 +oC = [-2.5, 2.5, -2.5, 2.5] +l = 1.49e-3 # wavelength +lList = np.linspace(1.49e-3, 1.78e-3, 11).tolist() + +dispMax = 17.9e-3 / 3600 # degrees +dispersion = np.linspace(0, dispMax, len(lList)).tolist() + +Exyz_list = [] + +print ('Your resolution is {}x{} with a pixel size of {}'.format(n,m,dpix)) +# Start of the ELT-lens system +# Initiate Fourier Optics, profile and wavefront error objects +for l in lList: + i = lList.index(l) + print(l, dispersion[i]) + fou = FourierOptics(l=l) + prof = Profiles() + wfe = WavefrontError(n,m) + + + # A = prof.top_hat(38854/2, n, m, dpix=dpix, oversample=1) + A = prof.ELT_pupil(n,m,dpix) + + # 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 + + z = wfe.generate_sine_pattern(1e-4, 1e-4, 180, 180) + dispersion_tilt = wfe.generate_tilt(x,y,tx=0, ty=dispersion[i]) + # psd_surface = wfe.generate_height_map(30e-6, 2, 1) + + M1.add_wavefront_error(z+dispersion_tilt)# + psd_surface) + t_l = M1.phase_delay(x,y,l) + U_in2 = A * t_l + + + # Propagating the original, unperturbed, system and calculating encircled energy + Exyz_s,xs,ys = fou.fresnel_propagation(U_in2, x, y, dz=M1.f, show=False, outputCoordinates=oC) + + Exyz_list.append(Exyz_s) + +Exyz_s = np.asarray(Exyz_list).sum(axis=0) +plt.imshow(np.abs(Exyz_s), cmap='hot');plt.colorbar();plt.show() +plt.imshow(np.abs(Exyz_s)**2, cmap='hot', norm=LogNorm(vmin=1e-7));plt.colorbar();plt.show() \ No newline at end of file diff --git a/scripts/high_contrast_imager.py b/scripts/high_contrast_imager.py new file mode 100644 index 0000000..ca7ba15 --- /dev/null +++ b/scripts/high_contrast_imager.py @@ -0,0 +1,184 @@ +#!/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 +from fourierOptics import * + + +def keplerIII_period_to_semimajor_axis(orbital_period): + ######################################################### + # Units: orbital period [yr], separation [au] # + # Function from + # https://github.com/katiebreivik/Keplers_Laws/blob/master/Graphing%20Kepler's%20Laws%20with%20Python.ipynb + ######################################################### + semimajor_axis_cubed = orbital_period**2 + semimajor_axis = semimajor_axis_cubed**(1./3.) + + return semimajor_axis + +def make_kepler_orbit(eccentricity,orbital_period, nStep=500): + ########################################################## + # Units: orbital period [years] # + # returns: 500 true anomaly values throughout the orbit # + # Function from + # https://github.com/katiebreivik/Keplers_Laws/blob/master/Graphing%20Kepler's%20Laws%20with%20Python.ipynb + ########################################################## + tRange=np.linspace(0.0,orbital_period,nStep) + theta=[] + for time in tRange: + PsiDiff=1.0 + M=2*np.pi*time/orbital_period + PsiOld=M + theta0old=180.0 + while PsiDiff>1e-10: + PsiNew=M+eccentricity*np.sin(PsiOld) + PsiDiff=PsiNew-PsiOld + PsiOld=PsiNew + theta0=2*np.arctan(((1+eccentricity)/(1-eccentricity))**(0.5)*np.tan(PsiOld/2.)) + theta.append(theta0) + return theta + +def orbit(semimajor_axis,eccentricity,true_anomaly): + ############################################## + # Units: separation [au] # + # Function from + # https://github.com/katiebreivik/Keplers_Laws/blob/master/Graphing%20Kepler's%20Laws%20with%20Python.ipynb + ############################################## + + # define the shape equation + r_orbit = semimajor_axis*(1 - eccentricity**2)/(1 + eccentricity*np.cos(true_anomaly)) + x_orbit = r_orbit*np.cos(true_anomaly) + y_orbit = r_orbit*np.sin(true_anomaly) + return x_orbit,y_orbit + + +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 = 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)) + + # Initiate the profiles class and pick a profile (tophat, obscured tophat or spider) + prof = Profiles() + + ############### + # from https://github.com/katiebreivik/Keplers_Laws/blob/master/Graphing%20Kepler's%20Laws%20with%20Python.ipynb + orbital_period = 1.0 + eccentricity = 0.23 + list1 = make_kepler_orbit(eccentricity,orbital_period, 360) + semimajor_axis = keplerIII_period_to_semimajor_axis(orbital_period) + xOrbit, yOrbit = orbit(semimajor_axis,eccentricity,list1) + ############### + + xOrbit *= 8 + yOrbit *= 8 + + # Create lens system + L1 = LensObject('N-BK7', -50, 1e99, 1, 15, name='L1') + L2 = LensObject('N-BK7', 50, 1e99, 1, 15, name='L2') + L3 = LensObject('N-BK7', 50, 1e99, 1, 15, name='L3') + L4 = LensObject('N-BK7', 50, 1e99, 1, 15, name='L4') + + # Calculate corresponding phase delays + t_l1 = L1.phase_delay(x,y,l, paraxial=True) + t_l2 = L2.phase_delay(x,y,l, paraxial=True) + t_l3 = L3.phase_delay(x,y,l, paraxial=True) + t_l4 = L4.phase_delay(x,y,l, paraxial=True) + + # Verify that the focal lengths are correct + print('Focal length of lens {} is {:.1f} mm.'.format(L1.name, L1.f)) + print('Focal length of lens {} is {:.1f} mm.'.format(L2.name, L2.f)) + print('Focal length of lens {} is {:.1f} mm.'.format(L3.name, L3.f)) + print('Focal length of lens {} is {:.1f} mm.'.format(L4.name, L4.f)) + + # Other parameters + oC = [x.min(), x.max(), y.min(), y.max()] + savename = False + + # The coronographic mask + coroMask = np.ones((n,m)) - prof.top_hat(2, n, m, dpix=dpix, dx=0, dy=0) + coroMask = np.ones((n,m)) + + # The Lyot stop + lyotStop = prof.obscured_top_hat(15, 3, n, m, dpix=dpix) + lyotStop = np.ones((n,m)) + + for i in range(len(list1)): + title = f"{i+1:03d}" + print (' Calculating {0}/{1}'.format(title, f"{len(list1):03d}"), end="\r") + intenseSource = prof.top_hat(1, n, m, dpix=dpix, oversample=1) + weakCompanion = prof.top_hat(0.1, n, m, dpix=dpix, oversample=1, dx=xOrbit[i], dy=yOrbit[i]) + A = intenseSource + 0.01*weakCompanion + A[A == 0] = 1e-99 + # plt.figure() + # plt.imshow(np.abs(A)**2, cmap='hot', extent=[x.min(), x.max(), y.min(), y.max()], norm=LogNorm(vmin=1e-6, vmax=1)) + # plt.xlim(-20, 20) + # plt.ylim(-20, 20) + # plt.colorbar() + # plt.savefig('./figures/high_contrast_imaging/input_{:02d}.png'.format(i+1), bbox_inches='tight', transparent=True) + # plt.close() + + # The input distribution gets a curved wavefront after the lens + U_in1 = A * t_l1 + + # Performing the propagation using Fresnel approximation to the focal plane. + Exyz,_,_ = fou.fresnel_propagation(U_in1, x, y, dz=-L1.f, show=False, save=savename, outputCoordinates=oC) + + # Focussing the beam again + U_in2 = Exyz * t_l2 * t_l2 + Exyz2,_,_ = fou.fresnel_propagation(U_in2, x, y, dz=L2.f, show=False, save=savename, outputCoordinates=oC) + + # Now we add a coronographic mask, following a gaussian + U_in3 = Exyz2.real * coroMask + Exyz2.imag * coroMask * 1j + + # Let this propagate again + Exyz3,_,_ = fou.fresnel_propagation(U_in3, x, y, dz=L2.f, show=False, save=savename, outputCoordinates=oC) + + # Collimating the beam + U_in4 = Exyz3 * t_l3 + + # Addition of a lyot stop + U_in4 = U_in4.real * lyotStop + U_in4.imag * lyotStop * 1j + + # Making an image + U_in5 = U_in4 * t_l4 + Exyz5,_,_ = fou.fresnel_propagation(U_in5, x, y, dz=L4.f, show=False, save=savename, outputCoordinates=oC) + # imageMask = np.ones((n,m)) - prof.top_hat(5.7, n, m, dpix=dpix, dx=0, dy=0) + # imageMask[ imageMask == 0] = 1e-99 + # imageData = np.abs(Exyz5.real * imageMask + Exyz5.imag * imageMask * 1j)**2 + imageData = np.abs(Exyz5)**2 + + plt.figure() + plt.imshow(imageData, cmap='hot', extent=[x.min(), x.max(), y.min(), y.max()], norm=LogNorm(vmin=1e-11)) + plt.xlim(-20, 20) + plt.ylim(-20, 20) + plt.colorbar() + plt.savefig('./figures/high_contrast_imaging/no_mask_image_{:03d}.png'.format(i+1), bbox_inches='tight', transparent=True) + # plt.show() + plt.close() + + + # End time counter + end = time.time() + + print('The program took {:.3f} seconds!'.format((end-start))) + + diff --git a/scripts/movie_atm_dispersion.py b/scripts/movie_atm_dispersion.py new file mode 100644 index 0000000..52fb935 --- /dev/null +++ b/scripts/movie_atm_dispersion.py @@ -0,0 +1,289 @@ +#!/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') +sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/atmosphere') +from fourierOptics import * +from atmosphericModel import * +from scipy import interpolate + +def perform_padding(A, xs, ys): + assert A.shape[0] <= xs and A.shape[1] <= ys, 'New size is smaller than array size.' + x_margin = int((xs - A.shape[0])/2) + y_margin = int((ys - A.shape[1])/2) + return np.pad(A, (x_margin, y_margin)) + +def find_nearest(array, value): + array = np.asarray(array) + diff = (np.abs(array - value)).argmin() + idx = np.unravel_index(diff, array.shape) + return idx + +def find_maximum(array): + idx = np.unravel_index(array.argmax(), array.shape) + return idx + +def select_subsection(x, y , outputCoordinates): + select_1 = find_nearest(x, outputCoordinates[0])[1] + select_2 = find_nearest(x, outputCoordinates[1])[1] + select_3 = find_nearest(y, outputCoordinates[2])[0] + select_4 = find_nearest(y, outputCoordinates[3])[0] + return [select_1, select_2, select_3, select_4] + +def perform_extrapolation(x, y, n, m): + # Determine properties of x,y meshgrid + xshape = x.shape[0] + yshape = y.shape[0] + dpix = (np.abs(np.min(x)) + np.abs(np.max(x))) / (xshape-1) + + # Extrapolate from, and include, the old x,y meshgrid + xnew = np.linspace(-(n-1)/2*dpix, (n-1)/2*dpix, n) + ynew = np.linspace(-(m-1)/2*dpix, (m-1)/2*dpix, m) + xnew, ynew = np.meshgrid(xnew, ynew) + return xnew, ynew + +def calculate_dispersion(n_iter): + # Parameters at Cerro Armazones + T = 279.65 # k + P = 71200 # Pa + H = 0.22 + xc = 450 # ppm + lat = -24.5983 # degrees + h = 3064 # m + l1 = 1.49 # micron + l2 = 1.78 # micron + + # Initializing dispersion model + at = Observatory() + + # Calculating indices of refraction for l1 and l2 + n1 = at.n_tph(l=l1, T=T, p=P, RH=H, xc=xc) + n2 = at.n_tph(l=l2, T=T, p=P, RH=H, xc=xc) + + # Density of the atmosphere (following CIPM-81/91 equations) + rho = at.rho(p=P, T=T, RH=H, xc=xc) + + # Initializing dispersion object + disp = dispersion(lat, h) + disp.setReducedHeight(P, rho) + + zenith_obs = np.linspace(0,75,n_iter) + dispersion_obs = [disp.cassini(n1, n2, z) for z in zenith_obs] + return dispersion_obs, zenith_obs + +if __name__ == '__main__': + plt.rc('font', family='serif') + plt.rc('axes', labelsize='large') + + + # Dispersion levels: + lList = np.linspace(1.49e-3, 1.78e-3, 30).tolist() + dispMaxList, zenith_obs = calculate_dispersion(25) + for dispMax in dispMaxList: + dispersion = np.linspace(0, dispMax, len(lList)).tolist() + zenith_angle = zenith_obs[dispMaxList.index(dispMax)] + print(zenith_angle) + Exyz_list = [] + + # Input parameters: + dpix = 50 + n, m = 3001, 3001 + x, y = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m)) + print ('Your resolution is {}x{} with a pixel size of {}'.format(n,m,dpix)) + print('This means your grid is {} mm by {} mm, with each pixel {} mm.'.format(x.max()-x.min(), y.max()-y.min(), dpix)) + + ########################################################### + #################### 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) + + # plt.imshow(ELT_MASK);plt.colorbar();plt.show() + + + + + + + + ########################################################### + + for l in lList: + i = lList.index(l) + print('Wavelength is {:.6f} mm. ({:02d}/{}).'.format(l,lList.index(l)+1, len(lList))) + x, y = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m)) + + fou = FourierOptics(l=l) + wfe = WavefrontError(n,m) + dispersion_tilt = wfe.generate_tilt(x, y, tx=0, ty=dispersion[i]) + + # ELT Mirror system combined into a single lens + ELT_MIRROR = ParabolicMirror(dx=0, dy=0, name='ELT-paraxial', CA=38542, f=6.84022e5) + + ELT_FP = OpticalSurface(name='MCD-IFP') + + ELT_COL = ParaxialLens(CA=154.54, f=1418.94, name='MCD-COL') + + ELT_ADC = ADCObject() + + # optAngle = ELT_ADC.find_optimum_angle(17.9, lList[0], lList[-1]) + # print(optAngle) + # print('ADC angle now set to {:.4f}'.format(ELT_ADC.theta)) + # ELT_ADC.set_orientation(180) + + 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() + + ########################################################### + ###################### 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], + [-3.5, 3.5, -3.5, 3.5] + ] + + xin, yin = x, y + + print('Starting propagation...') + + # From ELT-M1 to ELT-FP + print('Propagating from ELT-M1 to ELT-IFP', end='\r') + wf = ELT_MASK.astype('complex128') + ELT_MIRROR.add_wavefront_error(dispersion_tilt) + 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', end='\r') + 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', end='\r') + 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', end='\r') + 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', end='\r') + 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', end='\r') + 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', end='\r') + 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', end='\r') + 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=coordinatesList[-1], apply_interpolation=True) + + # 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...', end='\r') + # 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!') + + + Exyz_list.append(np.abs(Exyz8)**2) + fou = None + del fou + wfe = None + del wfe + MCD = None + del MCD + wf = None + del wf + ELT_MIRROR = None + del ELT_MIRROR + ELT_COL = None + del ELT_COL + ELT_ADC = None + del ELT_ADC + ELT_CAM = None + del ELT_CAM + ELT_DET = None + del ELT_DET + ########################################################### + + + Exyz_total = np.asarray(Exyz_list).sum(axis=0) + if zenith_angle == zenith_obs[0]: + maxIntensity = np.max(Exyz_total) + + _xip = x8 + _yip = y8 + save = 'ELT-MCD-Hband_dispersion_z{:02d}.pdf'.format(int(zenith_angle)) + _fig = plt.figure() + _ax = _fig.add_subplot(111) + _im = _ax.imshow(Exyz_total.T, cmap='hot', + extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()], + norm=LogNorm(vmax=maxIntensity, vmin=2e-7)) + _ax.set_xlim([_xip.min(), _xip.max()]) + _ax.set_ylim([_yip.min(), _yip.max()]) + _ax.set_xlabel('x [mm]') + _ax.set_ylabel('y [mm]') + _fig.colorbar(_im, ax=_ax, label='Intensity') + plt.tight_layout() + plt.savefig(save, bbox_inches='tight', transparent=True) + plt.close()#show() + + set_trace() + + + + + diff --git a/scripts/off_axis_test.py b/scripts/off_axis_test.py new file mode 100644 index 0000000..36524dc --- /dev/null +++ b/scripts/off_axis_test.py @@ -0,0 +1,67 @@ +#!/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 * + +def phaseTilt(x,y,tx,ty): + phaseTilt = x * np.tan(np.deg2rad(tx)) + y * np.tan(np.deg2rad(ty)) + return phaseTilt + +if __name__ == '__main__': + plt.rc('font', family='serif') + plt.rc('axes', labelsize='large') + + # Start time counter + start = time.time() + + lList = [30e-3, 32.5e-3, 35e-3, 37.5e-3, 40e-3] + Exyz_list = [] + # Pick a wavelength + l = 30e-3 + + for l in lList: + # Initiate Fourier Optics + fou = FourierOptics(l=l) + + # Pick sampling parameters + dpix = 1/20 + 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) + # A = prof.ELT_pupil(n,m,dpix) + plt.imshow(np.abs(A)) + + # 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)) + + + # Create a lens with a focal length ~100 mm, with a clear aperture radius of 4 mm and thickness of 1 mm. + L1 = SphericalLens('N-BK7', thickness=5, R1=100, R2=300, CA=5, name='L1') + print('Focal length of lens {} is {:.1f} mm.'.format(L1.name, L1.f)) + + # L1.add_wavefront_error(phaseTilt(x,y,tx=0,ty=2.5)) + + # 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()] + # oC = [0,x.max(),0,y.max()] + savename = False + Exyz_og,_,_ = fou.fresnel_propagation(U_in, x, y, dz=L1.f, show=True, save=savename, outputCoordinates=oC) + Exyz_list.append(Exyz_og) + + Exyz = np.asarray(Exyz_list).sum(axis=0) + plt.imshow(np.abs(Exyz), cmap='hot');plt.colorbar();plt.show() + + diff --git a/scripts/planewave_rotation_test.py b/scripts/planewave_rotation_test.py new file mode 100644 index 0000000..a3c91e8 --- /dev/null +++ b/scripts/planewave_rotation_test.py @@ -0,0 +1,460 @@ +#!/usr/bin/env Python3 +import numpy as np +from pdb import set_trace +import matplotlib.pyplot as plt +from matplotlib import cm, colors +from cmath import sqrt as csqrt +from scipy import interpolate +from scipy.optimize import leastsq +from matplotlib.ticker import FormatStrFormatter +from matplotlib.colors import LogNorm +import pyfftw +# import glassCatalog +from tabulate import tabulate +import types +import time +from math import ceil +import sys +from mpl_toolkits import mplot3d +sys.path.insert(0, '../') +import fourierOptics as fouOpt + + + +class FourierOptics: + def __init__(self, l): + self.l = l # Wavelength + + + def plane_wave_propagation_rotation(self, A, dpix=None, Rx=None, Ry=None, oversample=None, show=False, save=False, scaleInputArray=False, outputLimits=False): + ''' + Propagation through the Fresnel approximation, which is valid in many different situations. + A separate code for the Fraunhofer propagation is not striclty necessary, as the exponent + converges to 1 for large z values in any case. + + Input parameters: + A - [array] (Complex) Input field + dpix - [float] pixel scale at which the calculations are done + dz - [float] propagation distance + oversample - [bool] Apply padding on the outside of A, increasing the size of A + by the given factor. + show - [bool] Show the resulting amplitude and phase map after propagation + save - [bool] Save the figure + scaleInputArray - [tuple] Rescale the resolution of the input array using interpolation. + outputCoordinates - [list] list the corners of the desired output coordinate plane. + Order: [xmin, xmax, ymin, ymax] + + Output: + E_out - [array] (Complex) Output field distribution + x, y - [array] Output coordinates + ''' + + if scaleInputArray: + A = np.asarray(A) + _n, _m = A.shape + _nA, _mA = np.linspace(-1, 1, _n), np.linspace(-1, 1, _m) + A_ip_real = interpolate.interp2d(_nA, _mA, A.real, kind='cubic') + A_ip_imag = interpolate.interp2d(_nA, _mA, A.imag, kind='linear') + + _a, _b = scaleInputArray + _nx, _my = np.linspace(-1, 1, _a), np.linspace(-1, 1, _b) + A_out_ip_real = A_ip_real(_nx, _my) + A_out_ip_imag = A_ip_imag(_nx, _my) + A_out_ip = A_out_ip_real + A_out_ip_imag*1j + A = np.rint(A_out_ip_real) + + if oversample: + _Aos = np.zeros((np.shape(A)[0]*oversample, np.shape(A)[1]*oversample)) + _Aos[int(len(_Aos)/2-len(A)/2): int(len(_Aos)/2+len(A)/2), int(len(_Aos)/2-len(A)/2): int(len(_Aos)/2+len(A)/2)] = A + A = _Aos + + _n, _m = np.shape(A) + + # A at z=0 is then + _E_xy0 = pyfftw.empty_aligned((_n,_m), dtype='complex128') + _E_xy0[:] = np.fft.ifftshift(A) # use ifftshift to shift origin to start array + _IntensityNormFactor = np.sum(np.abs(_E_xy0)**2) # Factor to scale the intensity after the FFTs + _Az0 = pyfftw.empty_aligned((_n,_m), dtype='complex128') + _FFTW_Obj = pyfftw.FFTW( _E_xy0, _Az0, axes=(0,1), direction='FFTW_BACKWARD', threads=6, flags=('FFTW_ESTIMATE',)) + _FFTW_Obj() + + # E-field at z=dz is the fourier transform of A(z=0)*e^(-ikz) + _fx_l = np.fft.fftfreq(_n, d=dpix) + _fy_l = np.fft.fftfreq(_m, d=dpix) + _fx, _fy = np.meshgrid(_fx_l, _fy_l) + + # # Plotting stuff + # if False: + # Rx0 = Rx + # Ry0 = Ry + + # # Figure 3 from Matsushima et al., 2003 + # fig = plt.figure() + # ax = fig.add_subplot(121) + # ax.grid(alpha=0.3) + # ax.set_xlabel('u') + # ax.set_ylabel('$\\hat{u}$') + # ax.set_aspect('equal') + # ax2 = fig.add_subplot(122) + # ax2.grid(alpha=0.3) + # ax2.set_xlabel('v') + # ax2.set_ylabel('$\\hat{v}$') + # ax2.set_aspect('equal') + + # for theta in [0]:#,30,45,60,90]: + # Rx = 60 + # Ry = 60 + # print(Rx, Ry) + # Rx = np.deg2rad(Rx) + # Ry = np.deg2rad(Ry) + + # _Rx = np.array([[1,0,0], + # [0, np.cos(Rx), -np.sin(Rx)], + # [0, np.sin(Rx), np.cos(Rx)]]) + # _Ry = np.array([[np.cos(Ry), 0, np.sin(Ry)], + # [0,1,0], + # [-np.sin(Ry), 0, np.cos(Ry)]]) + # _R = np.dot(_Ry,_Rx) + # print(_R) + # _fz = np.sqrt( (self.l**(-2) - _fx**2 - _fy**2).astype('complex128') ) + # _fxr = (_R[0,0]*_fx + _R[1,0]*_fy + _R[2,0]*_fz) + # _fyr = (_R[0,1]*_fx + _R[1,1]*_fy + _R[2,1]*_fz) + # _fzr = np.sqrt( (self.l**(-2) - _fxr**2 - _fyr**2).astype('complex128') ) + + # _u0 = _R[2,0]*1/self.l + # _v0 = _R[2,1]*1/self.l + # # Recenter the new rotated frequencies at 0,0 + # _fxr_c = (_fxr - _u0) + # _fyr_c = (_fyr - _v0) + # _fzr_c = np.sqrt( (self.l**(-2) - _fxr_c**2 - _fyr_c**2).astype('complex128') ) + + # # _Az02 = np.fft.fftshift(_Az0) + # _fx_l2 = np.fft.fftshift(_fx_l) + # _fy_l2 = np.fft.fftshift(_fy_l) + # _fxr2 = np.fft.fftshift(_fxr) + # _fyr2 = np.fft.fftshift(_fyr) + # _fxr_c2 = np.fft.fftshift(_fxr_c) + # _fyr_c2 = np.fft.fftshift(_fyr_c) + # ax.plot(_fx_l2, _fxr2[1000], label='$\\phi_y={:.0f}$'.format(np.degrees(Ry))) + # ax2.plot(_fy_l2, _fyr2[:,1000], label='$\\phi_x={:.0f}$'.format(np.degrees(Rx))) + # ax.legend() + # ax2.legend() + # ax.set_xlim([-50,50]) + # ax.set_ylim([-50,50]) + # ax2.set_xlim([-50,50]) + # ax2.set_ylim([-50,50]) + # plt.tight_layout() + # plt.show() + + # Rx = Rx0 + # Ry = Ry0 + + # Figure showing frequency shift after Fourier rotatation + # fig = plt.figure(figsize=(8,3.5)) + # ax = fig.add_subplot(111) + # ax.plot(_fyr2[:,1000], np.abs(_Az02)[:,1000], 'k-', label='Reference frame ($\\hat{u}$)') + # ax.plot(_fy_l2, np.abs(_Az02)[:,1000], 'k-', alpha=0.4, label='Source frame ($u$)') + # ax.set_xlim([-0.5,+0.5]) + # ax.set_ylim(bottom=0) + # ax.set_xlabel('Frequency') + # ax.set_ylabel('Amplitude') + # ax.axes.get_xaxis().set_ticks([_v0,0]) + # ax.axes.xaxis.set_ticklabels(['$\\hat{u}=u_0$', 0]) + # ax.axes.get_yaxis().set_ticks([]) + # plt.grid(alpha=0.2) + # plt.legend() + # plt.savefig('frequency_shift_illustration.pdf', bbox_inches='tight', transparent=True) + # plt.show() + + + + + Rx = np.deg2rad(Rx) + Ry = np.deg2rad(Ry) + + _Rx = np.array([[1,0,0], + [0, np.cos(Rx), -np.sin(Rx)], + [0, np.sin(Rx), np.cos(Rx)]]) + _Ry = np.array([[np.cos(Ry), 0, np.sin(Ry)], + [0,1,0], + [-np.sin(Ry), 0, np.cos(Ry)]]) + _R = np.dot(_Rx,_Ry) + _fz = np.sqrt( (self.l**(-2) - _fx**2 - _fy**2).astype('complex128') ) + _fxr = (_R[0,0]*_fx + _R[1,0]*_fy + _R[2,0]*_fz) + _fyr = (_R[0,1]*_fx + _R[1,1]*_fy + _R[2,1]*_fz) + _fzr = np.sqrt( (self.l**(-2) - _fxr**2 - _fyr**2).astype('complex128') ) + + _u0 = _R[2,0]*1/self.l + _v0 = _R[2,1]*1/self.l + # Recenter the new rotated frequencies at 0,0 + _fxr_c = (_fxr - _u0) + _fyr_c = (_fyr - _v0) + _fzr_c = np.sqrt( (self.l**(-2) - _fxr_c**2 - _fyr_c**2).astype('complex128') ) + + set_trace() + # Calculation of the Jacobian + _J = (_R[0,0]*_R[1,1] - _R[0,1]*_R[1,0]) + (_R[0,1]*_R[1,2] - _R[0,2]*_R[1,1])*_fxr_c/_fzr_c + (_R[0,2]*_R[1,0]-_R[0,0]*_R[1,2])*_fyr_c/_fzr_c + # _J = (_R[0,0]*_R[1,1] - _R[0,1]*_R[1,0]) + (_R[1,0]*_R[2,0] - _R[2,0]*_R[1,1])*_fxr_c/_fzr_c + (_R[0,2]*_R[2,0]-_R[0,0]*_R[2,1])*_fyr_c/_fzr_c + + _E_rot = pyfftw.empty_aligned((_n,_m), dtype='complex128') + _A_in = pyfftw.empty_aligned((_n,_m), dtype='complex128') + _A_in[:] = _Az0 * np.abs(_J) + _FFTW_Obj2 = pyfftw.FFTW( _A_in, _E_rot, axes=(0,1), direction='FFTW_BACKWARD', threads=6, flags=('FFTW_ESTIMATE',)) + _FFTW_Obj2() + + # Coordinates of E_xy0 + _y, _x = np.mgrid[-_n / 2 * dpix : _n / 2 * dpix : 1j*_n, -_m / 2 * dpix : _m / 2 * dpix : 1j * _m] + + _carr_freq = np.exp(1j * 2 * np.pi * (_u0 * _x + _v0 * _y)) + _E_rot = np.fft.fftshift(_E_rot) # use fftshift to move the origin from edge to center + _E_rot *= _carr_freq + _E_rot /= np.sqrt(_IntensityNormFactor) # Normalization back to initial energy. + + + if Rx or Ry: + # A rotation requires interpolation over the output + # Performing interpolation to get to the required final plane, done for real and imaginary part. + _xls = np.linspace(-_n/2*dpix, _n/2*dpix, _n) + _yls = np.linspace(-_m/2*dpix, _m/2*dpix, _m) + _zls = np.zeros(_xls.shape) + _x = _R[0,0]*_xls + _R[1,0]*_yls + _R[2,0]*_zls # Rotated original coordinates + _y = _R[0,1]*_xls + _R[1,1]*_yls + _R[2,1]*_zls # Rotated original coordinates + _z = _R[0,2]*_xls + _R[1,2]*_yls + _R[2,2]*_zls + + _ys, _xs = np.mgrid[-_n / 2 * dpix : _n / 2 * dpix : 1j*_n, -_m / 2 * dpix : _m / 2 * dpix : 1j * _m] + # _xs = np.linspace(-_n/2*dpix, _n/2*dpix, _n) + # _ys = np.linspace(-_m/2*dpix, _m/2*dpix, _m) + _zs = np.zeros(_xs.shape) + + _xr = _R[0,0]*_xs + _R[1,0]*_ys + _R[2,0]*_zs # Rotated original coordinates + _yr = _R[0,1]*_xs + _R[1,1]*_ys + _R[2,1]*_zs # Rotated original coordinates + _zr = _R[0,2]*_xs + _R[1,2]*_ys + _R[2,2]*_zs + + + + # fc = np.fft.fftshift(np.angle(_E_xy0))[::3,::3] + # fcr = np.angle(_E_rot)[::3,::3] + # norm=colors.Normalize() + + # fig = plt.figure(figsize=(8,8)) + # ax = plt.axes(projection='3d') + # # ax.plot_surface(_zs[::3,::3], _xs[::3,::3], _ys[::3,::3], facecolors=cm.hot(norm(fc)), alpha=0.5) + # ax.plot_surface(_zr[::3,::3], _xr[::3,::3], _yr[::3,::3], facecolors=cm.hot(norm(fcr))) + # ax.set_xlabel('z') + # ax.set_ylabel('x') + # ax.set_zlabel('y') + # ax.set_xlim([-2,2]) + # ax.set_ylim([-2,2]) + # ax.set_zlim([-2,2]) + # plt.show() + + + _xs = np.linspace(-_n/2*dpix, _n/2*dpix, _n) + _ys = np.linspace(-_m/2*dpix, _m/2*dpix, _m) + _zs = np.zeros(_xs.shape) + + _xr = _R[0,0]*_xs + _R[1,0]*_ys + _R[2,0]*_zs # Rotated original coordinates + _yr = _R[0,1]*_xs + _R[1,1]*_ys + _R[2,1]*_zs # Rotated original coordinates + # _zr = _R[0,2]*_xs + _R[1,2]*_ys + _R[2,2]*_zs + # set_trace() + + _f_ip_real = interpolate.interp2d(_xr, _yr, _E_rot.real, kind='linear') + _f_ip_imag = interpolate.interp2d(_xr, _yr, _E_rot.imag, kind='linear') + _xi = np.linspace(_xr.min(), _xr.max(), _n) + _yi = np.linspace(_yr.min(), _yr.max(), _m) + _xi = _xs + _yi = _ys + _E_out_ip_real = _f_ip_real(_xi, _yi) + _E_out_ip_imag = _f_ip_imag(_xi, _yi) + _E_rot = _E_out_ip_real + _E_out_ip_imag*1j + _y, _x = np.meshgrid(_yi, _xi) + # set_trace() + + + + + + if show or save: + _E_rot_extent = [np.min(_x), np.max(_x), np.min(_y), np.max(_x)] + fig = plt.figure(figsize=(8,3)) + ax = fig.add_subplot(121) + im = ax.imshow(np.abs(_E_rot), cmap='hot', extent=_E_rot_extent)#, vmax=1, vmin=0) + ax.set_xlabel('x [mm]') + ax.set_ylabel('y [mm]') + fig.colorbar(im, ax=ax, label='Normalized intensity') + ax2 = fig.add_subplot(122) + im2 = ax2.imshow(np.angle(_E_rot), cmap='hot', extent=_E_rot_extent) + ax2.set_xlabel('x [mm]') + ax2.set_ylabel('y [mm]') + ax.set_xlim([-0.4,0.4]) + ax.set_ylim([-0.4,0.4]) + ax2.set_xlim([-0.4,0.4]) + ax2.set_ylim([-0.4,0.4]) + fig.colorbar(im2, ax=ax2) + plt.tight_layout() + if save: + plt.savefig(save, bbox_inches='tight', transparent=True) + if show: + plt.show() + elif not show: + plt.close() + + return np.array([_E_rot, _x, _y]) + + + def phase_exponent(self, fx, fy, l): + k = 2*np.pi / l + kx = np.asarray(2*np.pi * fx) + ky = np.asarray(2*np.pi * fy) + K_squared = (kx**2 + ky**2) + kA = np.ones(K_squared.shape) * k + K_1 = np.sqrt((kA**2-K_squared).astype('complex128')) + kz = K_1.real + -1j*K_1.imag + return kz + + + def coordinate_rotation(self, Rx, Ry, n=5, m=5): + _n = n + _m = m + Rx = np.deg2rad(Rx) + Ry = np.deg2rad(Ry) + dpix = 1/30 * self.l * 1000 + + _Rx = np.array([[1,0,0], + [0, np.cos(Rx), -np.sin(Rx)], + [0, np.sin(Rx), np.cos(Rx)]]) + _Ry = np.array([[np.cos(Ry), 0, np.sin(Ry)], + [0,1,0], + [-np.sin(Ry), 0, np.cos(Ry)]]) + _R = np.dot(_Rx,_Ry) + + theta = range(360) + _xcirc = 1.2 * np.cos(np.deg2rad(theta)) + _ycirc = 1.2 * np.sin(np.deg2rad(theta)) + _zcirc = np.zeros(_xcirc.shape) + + + _ys, _xs = np.mgrid[-_n / 2 * dpix : _n / 2 * dpix : 1j*_n, -_m / 2 * dpix : _m / 2 * dpix : 1j * _m] + _zs = np.zeros(_xs.shape) + + _xr = _R[0,0]*_xs + _R[1,0]*_ys + _R[2,0]*_zs # Rotated original coordinates + _yr = _R[0,1]*_xs + _R[1,1]*_ys + _R[2,1]*_zs # Rotated original coordinates + _zr = _R[0,2]*_xs + _R[1,2]*_ys + _R[2,2]*_zs + + _xcirc_r = _R[0,0]*_xcirc + _R[1,0]*_ycirc + _R[2,0]*_zcirc # Rotated original coordinates + _ycirc_r = _R[0,1]*_xcirc + _R[1,1]*_ycirc + _R[2,1]*_zcirc # Rotated original coordinates + _zcirc_r = _R[0,2]*_xcirc + _R[1,2]*_ycirc + _R[2,2]*_zcirc + + _xcirc_p = _xcirc # Projected original coordinates + _ycirc_p = _ycirc # Projected original coordinates + # _zcirc_p = _R[0,2]*_xcirc*np.sqrt(1+np.tan(Ry)**2) + _R[1,2]*_ycirc*np.sqrt(1+np.tan(Rx)**2) + _zcirc_p = np.tan(Ry)*_xcirc - np.tan(Rx)*_ycirc + _zcirc + # _zcirc_p = -_xcirc / np.tan(Ry) - _ycirc / np.tan(Rx) + _zcirc + + # Plot 3d Wireframe of source plane and reference plane + fig = plt.figure(figsize=(8,8)) + ax = plt.axes(projection='3d') + ax.plot_wireframe(_zs, _xs, _ys, label='source frame', color='blue', alpha=.1) + ax.plot_wireframe(_zr, _xr, _yr, label='reference frame', color='red', alpha=.1) + ax.plot3D(_zcirc, _xcirc, _ycirc, 'blue') + ax.plot3D(_zcirc_r, _xcirc_r, _ycirc_r, 'red', alpha=.2) + ax.plot3D(_zcirc_p, _xcirc_p, _ycirc_p, 'black') + ax.set_xlabel('z') + ax.set_ylabel('x') + ax.set_zlabel('y') + ax.set_xlim([-3,3]) + ax.set_ylim([-3,3]) + ax.set_zlim([-3,3]) + ax.grid(True, alpha=0.3) + ax.legend() + plt.tight_layout() + # plt.savefig('./coordinate_rotation_figures/coordinate_rotation_rx{:.0f}_ry{:.0f}.png'.format(np.degrees(Rx), np.degrees(Ry)), bbox_inches='tight') + # plt.close() + plt.show() + + + + +# _Az02 = np.fft.fftshift(_Az0) +# _fx_l2 = np.fft.fftshift(_fx_l) +# _fy_l2 = np.fft.fftshift(_fy_l) +# _fxr2 = np.fft.fftshift(_fxr) +# _fyr2 = np.fft.fftshift(_fyr) +# plt.plot(_fyr2[:,1000], np.abs(_Az02)[:,1000]); plt.plot(_fy_l2, np.abs(_Az02)[:,1000]); plt.grid(); plt.show() + +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) + fouOpt_norm = fouOpt.FourierOptics(l=l) + + # Pick sampling parameters + dpix = 1/30*l + n, m = 3001, 3001 + 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 = fouOpt.Profiles() + # A = prof.spiderTopHat(250, 100, 2001, 2001, dpix=1, oversample=1) + A = prof.top_hat(4*l, n, m, dpix=dpix, oversample=1) + # x, y = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m)) + # gauss = fouOpt.GaussianBeam(0.02, l) + # A = gauss.fundamental_mode(.4*l,n,m,np.max(x)) + # Exyz = A + + + # # Calculate the Gaussian beam for the whole tilted image plane + # A_at_z4l = gauss.fundamental_mode(4*l,n,m,np.max(x)) + + # The input distribution gets a curved wavefront after the lens + U_in = A + # fou.coordinate_rotation(45,30) + # sys.exit() + Exyz,x4l,y4l = fouOpt_norm.plane_wave_propagation(A, dz=4*l, dpix=dpix, oversample=None, show=False) + + # _E_res_abs = np.abs(A_at_z4l)/np.max(np.abs(A_at_z4l)) - np.abs(Exyz)/np.max(np.abs(Exyz)) + # _E_res_ang = np.angle(A_at_z4l)/np.max(np.angle(A_at_z4l)) - np.angle(Exyz)/np.max(np.angle(Exyz)) + + # _E_rot_extent = [np.min(x4l), np.max(x4l), np.min(y4l), np.max(x4l)] + # fig = plt.figure(figsize=(8,3)) + # ax = fig.add_subplot(121) + # im = ax.imshow(_E_res_abs.T, cmap='hot', extent=_E_rot_extent)#, vmax=1, vmin=0) + # ax.set_xlabel('x [mm]') + # ax.set_ylabel('y [mm]') + # fig.colorbar(im, ax=ax, label='Normalized intensity') + # ax2 = fig.add_subplot(122) + # im2 = ax2.imshow(_E_res_ang.T, cmap='hot', extent=_E_rot_extent) + # ax2.set_xlabel('x [mm]') + # ax2.set_ylabel('y [mm]') + # ax.set_xlim([-0.4,0.4]) + # ax.set_ylim([-0.4,0.4]) + # ax2.set_xlim([-0.4,0.4]) + # ax2.set_ylim([-0.4,0.4]) + # fig.colorbar(im2, ax=ax2) + # plt.tight_layout() + # # plt.show()/ + # plt.close() + + # Erot,_,_ = fou.plane_wave_propagation_rotation(Exyz, dpix=dpix, Rx=0, Ry=45, oversample=None, show=False) + # Erot2,_,_ = fou.plane_wave_propagation_rotation(Exyz, dpix=dpix, Rx=0, Ry=45, oversample=None, show=True) + Erot2,x,y = fou.plane_wave_propagation_rotation(Exyz, dpix=dpix, Rx=60, Ry=0, oversample=None, show=True)#, save='p2p_gauss_Rx45_Ry45.png') + + + + + + + + + + + + \ No newline at end of file diff --git a/scripts/planewave_rotation_test_2020.py b/scripts/planewave_rotation_test_2020.py new file mode 100644 index 0000000..12a90d9 --- /dev/null +++ b/scripts/planewave_rotation_test_2020.py @@ -0,0 +1,446 @@ +#!/usr/bin/env Python3 +import numpy as np +from pdb import set_trace +import matplotlib.pyplot as plt +from matplotlib import cm, colors +from cmath import sqrt as csqrt +from scipy import interpolate +from scipy.optimize import leastsq +from matplotlib.ticker import FormatStrFormatter +from matplotlib.colors import LogNorm +import pyfftw +# import glassCatalog +from tabulate import tabulate +import types +import time +from math import ceil +import sys +from mpl_toolkits import mplot3d +sys.path.insert(0, '../../') +import fourierOptics as fouOpt + + + +class FourierOptics: + def __init__(self, l): + self.l = l # Wavelength + + + def plane_wave_propagation_rotation(self, A, dpix=None, Rx=None, Ry=None, oversample=None, show=False, save=False, scaleInputArray=False, outputLimits=False): + ''' + Propagation through the Fresnel approximation, which is valid in many different situations. + A separate code for the Fraunhofer propagation is not striclty necessary, as the exponent + converges to 1 for large z values in any case. + + Input parameters: + A - [array] (Complex) Input field + dpix - [float] pixel scale at which the calculations are done + dz - [float] propagation distance + oversample - [bool] Apply padding on the outside of A, increasing the size of A + by the given factor. + show - [bool] Show the resulting amplitude and phase map after propagation + save - [bool] Save the figure + scaleInputArray - [tuple] Rescale the resolution of the input array using interpolation. + outputCoordinates - [list] list the corners of the desired output coordinate plane. + Order: [xmin, xmax, ymin, ymax] + + Output: + E_out - [array] (Complex) Output field distribution + x, y - [array] Output coordinates + ''' + + if scaleInputArray: + A = np.asarray(A) + _n, _m = A.shape + _nA, _mA = np.linspace(-1, 1, _n), np.linspace(-1, 1, _m) + A_ip_real = interpolate.interp2d(_nA, _mA, A.real, kind='cubic') + A_ip_imag = interpolate.interp2d(_nA, _mA, A.imag, kind='linear') + + _a, _b = scaleInputArray + _nx, _my = np.linspace(-1, 1, _a), np.linspace(-1, 1, _b) + A_out_ip_real = A_ip_real(_nx, _my) + A_out_ip_imag = A_ip_imag(_nx, _my) + A_out_ip = A_out_ip_real + A_out_ip_imag*1j + A = np.rint(A_out_ip_real) + + if oversample: + _Aos = np.zeros((np.shape(A)[0]*oversample, np.shape(A)[1]*oversample)) + _Aos[int(len(_Aos)/2-len(A)/2): int(len(_Aos)/2+len(A)/2), int(len(_Aos)/2-len(A)/2): int(len(_Aos)/2+len(A)/2)] = A + A = _Aos + + _n, _m = np.shape(A) + + # A at z=0 is then + _E_xy0 = pyfftw.empty_aligned((_n,_m), dtype='complex128') + _E_xy0[:] = np.fft.ifftshift(A) # use ifftshift to shift origin to start array + _IntensityNormFactor = np.sum(np.abs(_E_xy0)**2) # Factor to scale the intensity after the FFTs + _Az0 = pyfftw.empty_aligned((_n,_m), dtype='complex128') + _FFTW_Obj = pyfftw.FFTW( _E_xy0, _Az0, axes=(0,1), direction='FFTW_BACKWARD', threads=6, flags=('FFTW_ESTIMATE',)) + _FFTW_Obj() + + # E-field at z=dz is the fourier transform of A(z=0)*e^(-ikz) + _fx_l = np.fft.fftfreq(_n, d=dpix) + _fy_l = np.fft.fftfreq(_m, d=dpix) + _fx, _fy = np.meshgrid(_fx_l, _fy_l) + + # # Plotting stuff + # if False: + # Rx0 = Rx + # Ry0 = Ry + + # # Figure 3 from Matsushima et al., 2003 + # fig = plt.figure() + # ax = fig.add_subplot(121) + # ax.grid(alpha=0.3) + # ax.set_xlabel('u') + # ax.set_ylabel('$\\hat{u}$') + # ax.set_aspect('equal') + # ax2 = fig.add_subplot(122) + # ax2.grid(alpha=0.3) + # ax2.set_xlabel('v') + # ax2.set_ylabel('$\\hat{v}$') + # ax2.set_aspect('equal') + + # for theta in [0]:#,30,45,60,90]: + # Rx = 60 + # Ry = 60 + # print(Rx, Ry) + # Rx = np.deg2rad(Rx) + # Ry = np.deg2rad(Ry) + + # _Rx = np.array([[1,0,0], + # [0, np.cos(Rx), -np.sin(Rx)], + # [0, np.sin(Rx), np.cos(Rx)]]) + # _Ry = np.array([[np.cos(Ry), 0, np.sin(Ry)], + # [0,1,0], + # [-np.sin(Ry), 0, np.cos(Ry)]]) + # _R = np.dot(_Ry,_Rx) + # print(_R) + # _fz = np.sqrt( (self.l**(-2) - _fx**2 - _fy**2).astype('complex128') ) + # _fxr = (_R[0,0]*_fx + _R[1,0]*_fy + _R[2,0]*_fz) + # _fyr = (_R[0,1]*_fx + _R[1,1]*_fy + _R[2,1]*_fz) + # _fzr = np.sqrt( (self.l**(-2) - _fxr**2 - _fyr**2).astype('complex128') ) + + # _u0 = _R[2,0]*1/self.l + # _v0 = _R[2,1]*1/self.l + # # Recenter the new rotated frequencies at 0,0 + # _fxr_c = (_fxr - _u0) + # _fyr_c = (_fyr - _v0) + # _fzr_c = np.sqrt( (self.l**(-2) - _fxr_c**2 - _fyr_c**2).astype('complex128') ) + + # # _Az02 = np.fft.fftshift(_Az0) + # _fx_l2 = np.fft.fftshift(_fx_l) + # _fy_l2 = np.fft.fftshift(_fy_l) + # _fxr2 = np.fft.fftshift(_fxr) + # _fyr2 = np.fft.fftshift(_fyr) + # _fxr_c2 = np.fft.fftshift(_fxr_c) + # _fyr_c2 = np.fft.fftshift(_fyr_c) + # ax.plot(_fx_l2, _fxr2[1000], label='$\\phi_y={:.0f}$'.format(np.degrees(Ry))) + # ax2.plot(_fy_l2, _fyr2[:,1000], label='$\\phi_x={:.0f}$'.format(np.degrees(Rx))) + # ax.legend() + # ax2.legend() + # ax.set_xlim([-50,50]) + # ax.set_ylim([-50,50]) + # ax2.set_xlim([-50,50]) + # ax2.set_ylim([-50,50]) + # plt.tight_layout() + # plt.show() + + # Rx = Rx0 + # Ry = Ry0 + + # Figure showing frequency shift after Fourier rotatation + # fig = plt.figure(figsize=(8,3.5)) + # ax = fig.add_subplot(111) + # ax.plot(_fyr2[:,1000], np.abs(_Az02)[:,1000], 'k-', label='Reference frame ($\\hat{u}$)') + # ax.plot(_fy_l2, np.abs(_Az02)[:,1000], 'k-', alpha=0.4, label='Source frame ($u$)') + # ax.set_xlim([-0.5,+0.5]) + # ax.set_ylim(bottom=0) + # ax.set_xlabel('Frequency') + # ax.set_ylabel('Amplitude') + # ax.axes.get_xaxis().set_ticks([_v0,0]) + # ax.axes.xaxis.set_ticklabels(['$\\hat{u}=u_0$', 0]) + # ax.axes.get_yaxis().set_ticks([]) + # plt.grid(alpha=0.2) + # plt.legend() + # plt.savefig('frequency_shift_illustration.pdf', bbox_inches='tight', transparent=True) + # plt.show() + + + + + Rx = np.deg2rad(Rx) + Ry = np.deg2rad(Ry) + + _Rx = np.array([[1,0,0], + [0, np.cos(Rx), -np.sin(Rx)], + [0, np.sin(Rx), np.cos(Rx)]]) + _Ry = np.array([[np.cos(Ry), 0, np.sin(Ry)], + [0,1,0], + [-np.sin(Ry), 0, np.cos(Ry)]]) + _R = np.dot(_Rx,_Ry) + _fz = np.sqrt( (self.l**(-2) - _fx**2 - _fy**2).astype('complex128') ) + _fxr = (_R[0,0]*_fx + _R[1,0]*_fy + _R[2,0]*_fz) + _fyr = (_R[0,1]*_fx + _R[1,1]*_fy + _R[2,1]*_fz) + _fzr = np.sqrt( (self.l**(-2) - _fxr**2 - _fyr**2).astype('complex128') ) + + _u0 = _R[2,0]*1/self.l + _v0 = _R[2,1]*1/self.l + # Recenter the new rotated frequencies at 0,0 + _fxr_c = (_fxr - _u0) + _fyr_c = (_fyr - _v0) + _fzr_c = np.sqrt( (self.l**(-2) - _fxr_c**2 - _fyr_c**2).astype('complex128') ) + + # Calculation of the Jacobian + # _J = (_R[0,0]*_R[1,1] - _R[0,1]*_R[1,0]) + (_R[0,1]*_R[1,2] - _R[0,2]*_R[1,1])*_fxr_c/_fzr_c + (_R[0,2]*_R[1,0]-_R[0,0]*_R[1,2])*_fyr_c/_fzr_c + _J = (_R[0,0]*_R[1,1] - _R[0,1]*_R[1,0]) + (_R[1,0]*_R[2,0] - _R[2,0]*_R[1,1])*_fxr_c/_fzr_c + (_R[0,2]*_R[2,0]-_R[0,0]*_R[2,1])*_fyr_c/_fzr_c + + _E_rot = pyfftw.empty_aligned((_n,_m), dtype='complex128') + _A_in = pyfftw.empty_aligned((_n,_m), dtype='complex128') + _A_in[:] = _Az0 * np.abs(_J) + _FFTW_Obj2 = pyfftw.FFTW( _A_in, _E_rot, axes=(0,1), direction='FFTW_BACKWARD', threads=6, flags=('FFTW_ESTIMATE',)) + _FFTW_Obj2() + + # Coordinates of E_xy0 + _y, _x = np.mgrid[-_n / 2 * dpix : _n / 2 * dpix : 1j*_n, -_m / 2 * dpix : _m / 2 * dpix : 1j * _m] + + _carr_freq = np.exp(1j * 2 * np.pi * (_u0 * _x + _v0 * _y)) + _E_rot = np.fft.fftshift(_E_rot) # use fftshift to move the origin from edge to center + _E_rot *= _carr_freq + _E_rot /= np.sqrt(_IntensityNormFactor) # Normalization back to initial energy. + + + if Rx or Ry: + # A rotation requires interpolation over the output + # Performing interpolation to get to the required final plane, done for real and imaginary part. + _xs = np.linspace(-_n/2*dpix, _n/2*dpix, _n) + _ys = np.linspace(-_m/2*dpix, _m/2*dpix, _m) + _zs = np.zeros(_xs.shape) + + _xr = _R[0,0]*_xs + _R[1,0]*_ys + _R[2,0]*_zs # Rotated original coordinates + _yr = _R[0,1]*_xs + _R[1,1]*_ys + _R[2,1]*_zs # Rotated original coordinates + # _zr = _R[0,2]*_xs + _R[1,2]*_ys + _R[2,2]*_zs + _f_ip_real = interpolate.interp2d(_xs, _ys, _E_rot.real, kind='linear') + _f_ip_imag = interpolate.interp2d(_xs, _ys, _E_rot.imag, kind='linear') + # _f_ip_real = interpolate.interp2d(_xr, _yr, _E_rot.real, kind='linear') + # _f_ip_imag = interpolate.interp2d(_xr, _yr, _E_rot.imag, kind='linear') + _xi = np.linspace(_xr.min(), _xr.max(), _n) + _yi = np.linspace(_yr.min(), _yr.max(), _m) + + # _xi = _xs + # _yi = _ys + _E_out_ip_real = _f_ip_real(_xi, _yi) + _E_out_ip_imag = _f_ip_imag(_xi, _yi) + _E_rot = _E_out_ip_real + _E_out_ip_imag*1j + _y, _x = np.meshgrid(_yi, _xi) + + + + + + + if show or save: + _E_rot_extent = [np.min(_x), np.max(_x), np.min(_x), np.max(_x)] + fig = plt.figure(figsize=(8,3)) + ax = fig.add_subplot(121) + im = ax.imshow(np.abs(_E_rot), cmap='hot', extent=_E_rot_extent)#, vmax=1, vmin=0) + ax.set_xlabel('x [mm]') + ax.set_ylabel('y [mm]') + fig.colorbar(im, ax=ax, label='Normalized intensity') + ax2 = fig.add_subplot(122) + im2 = ax2.imshow(np.angle(_E_rot), cmap='hot', extent=_E_rot_extent) + ax2.set_xlabel('x [mm]') + ax2.set_ylabel('y [mm]') + ax.set_xlim([-0.4,0.4]) + ax.set_ylim([-0.4,0.4]) + ax2.set_xlim([-0.4,0.4]) + ax2.set_ylim([-0.4,0.4]) + fig.colorbar(im2, ax=ax2) + plt.tight_layout() + if save: + plt.savefig(save, bbox_inches='tight', transparent=True) + if show: + plt.show() + elif not show: + plt.close() + + return np.array([_E_rot, _x, _y]) + + + def phase_exponent(self, fx, fy, l): + k = 2*np.pi / l + kx = np.asarray(2*np.pi * fx) + ky = np.asarray(2*np.pi * fy) + K_squared = (kx**2 + ky**2) + kA = np.ones(K_squared.shape) * k + K_1 = np.sqrt((kA**2-K_squared).astype('complex128')) + kz = K_1.real + -1j*K_1.imag + return kz + + + def coordinate_rotation(self, Rx, Ry, n=5, m=5): + _n = n + _m = m + Rx = np.deg2rad(Rx) + Ry = np.deg2rad(Ry) + dpix = 1/30 * self.l * 1000 + + _Rx = np.array([[1,0,0], + [0, np.cos(Rx), -np.sin(Rx)], + [0, np.sin(Rx), np.cos(Rx)]]) + _Ry = np.array([[np.cos(Ry), 0, np.sin(Ry)], + [0,1,0], + [-np.sin(Ry), 0, np.cos(Ry)]]) + _R = np.dot(_Rx,_Ry) + + theta = range(360) + _xcirc = 1.2 * np.cos(np.deg2rad(theta)) + _ycirc = 1.2 * np.sin(np.deg2rad(theta)) + _zcirc = np.zeros(_xcirc.shape) + + + _ys, _xs = np.mgrid[-_n / 2 * dpix : _n / 2 * dpix : 1j*_n, -_m / 2 * dpix : _m / 2 * dpix : 1j * _m] + _zs = np.zeros(_xs.shape) + + _xr = _R[0,0]*_xs + _R[1,0]*_ys + _R[2,0]*_zs # Rotated original coordinates + _yr = _R[0,1]*_xs + _R[1,1]*_ys + _R[2,1]*_zs # Rotated original coordinates + _zr = _R[0,2]*_xs + _R[1,2]*_ys + _R[2,2]*_zs + + _xcirc_r = _R[0,0]*_xcirc + _R[1,0]*_ycirc + _R[2,0]*_zcirc # Rotated original coordinates + _ycirc_r = _R[0,1]*_xcirc + _R[1,1]*_ycirc + _R[2,1]*_zcirc # Rotated original coordinates + _zcirc_r = _R[0,2]*_xcirc + _R[1,2]*_ycirc + _R[2,2]*_zcirc + + _xcirc_p = _xcirc # Projected original coordinates + _ycirc_p = _ycirc # Projected original coordinates + _zcirc_p = (-_R[2,0]*_xcirc_p - _R[2,1]*_ycirc_p) / _R[2,2] + + # Plot 3d Wireframe of source plane and reference plane + fig = plt.figure(figsize=(8,8)) + ax = plt.axes(projection='3d') + ax.plot_wireframe(_zs, _xs, _ys, label='source frame', color='blue', alpha=.3) + ax.plot_wireframe(_zr, _xr, _yr, label='reference frame', color='red', alpha=.3) + ax.plot3D(_zcirc, _xcirc, _ycirc, 'blue') + ax.plot3D(_zcirc_r, _xcirc_r, _ycirc_r, 'red', alpha=.2) + ax.plot3D(_zcirc_p, _xcirc_p, _ycirc_p, 'black') + ax.set_xlabel('z') + ax.set_ylabel('x') + ax.set_zlabel('y') + ax.set_xlim([-3,3]) + ax.set_ylim([-3,3]) + ax.set_zlim([-3,3]) + ax.grid(True, alpha=0.3) + ax.legend() + plt.tight_layout() + # plt.savefig('./coordinate_rotation_figures/coordinate_rotation_rx{:.0f}_ry{:.0f}.png'.format(np.degrees(Rx), np.degrees(Ry)), bbox_inches='tight') + # plt.close() + plt.show(block=False) + + + + +# _Az02 = np.fft.fftshift(_Az0) +# _fx_l2 = np.fft.fftshift(_fx_l) +# _fy_l2 = np.fft.fftshift(_fy_l) +# _fxr2 = np.fft.fftshift(_fxr) +# _fyr2 = np.fft.fftshift(_fyr) +# plt.plot(_fyr2[:,1000], np.abs(_Az02)[:,1000]); plt.plot(_fy_l2, np.abs(_Az02)[:,1000]); plt.grid(); plt.show() + +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) + fouOpt_norm = fouOpt.FourierOptics(l=l) + + # Pick sampling parameters + dpix = 1/40*l + n, m = 3001, 3001 + 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 = fouOpt.Profiles() + # A = prof.spiderTopHat(250, 100, 2001, 2001, dpix=1, oversample=1) + # A = prof.top_hat(4*l, n, m, dpix=dpix, oversample=1) + x, y = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m)) + gauss = fouOpt.GaussianBeam(0.02, l) + + A = gauss.fundamental_mode(0,n,m,np.max(x)) + # z0 = n*dpix * np.tan(np.deg2rad(30)) + # A = gauss.fundamental_mode(z0,n,m,np.max(x)) + # Exyz = A + + + # # Calculate the Gaussian beam for the whole tilted image plane + # A_at_z4l = gauss.fundamental_mode(4*l,n,m,np.max(x)) + + # The input distribution gets a curved wavefront after the lens + U_in = A + + # sys.exit() + Exyz,x4l,y4l = fouOpt_norm.plane_wave_propagation(A, dz=4*l, dpix=dpix, oversample=None, show=True) + fou.coordinate_rotation(45,20) + # _E_res_abs = np.abs(A_at_z4l)/np.max(np.abs(A_at_z4l)) - np.abs(Exyz)/np.max(np.abs(Exyz)) + # _E_res_ang = np.angle(A_at_z4l)/np.max(np.angle(A_at_z4l)) - np.angle(Exyz)/np.max(np.angle(Exyz)) + + # _E_rot_extent = [np.min(x4l), np.max(x4l), np.min(y4l), np.max(x4l)] + # fig = plt.figure(figsize=(8,3)) + # ax = fig.add_subplot(121) + # im = ax.imshow(_E_res_abs.T, cmap='hot', extent=_E_rot_extent)#, vmax=1, vmin=0) + # ax.set_xlabel('x [mm]') + # ax.set_ylabel('y [mm]') + # fig.colorbar(im, ax=ax, label='Normalized intensity') + # ax2 = fig.add_subplot(122) + # im2 = ax2.imshow(_E_res_ang.T, cmap='hot', extent=_E_rot_extent) + # ax2.set_xlabel('x [mm]') + # ax2.set_ylabel('y [mm]') + # ax.set_xlim([-0.4,0.4]) + # ax.set_ylim([-0.4,0.4]) + # ax2.set_xlim([-0.4,0.4]) + # ax2.set_ylim([-0.4,0.4]) + # fig.colorbar(im2, ax=ax2) + # plt.tight_layout() + # # plt.show()/ + # plt.close() + + Erot2,x,y = fou.plane_wave_propagation_rotation(Exyz, dpix=dpix, Rx=45, Ry=20, oversample=None, show=True)#, save='p2p_gauss_Rx45_Ry45.png') + + ''' + # Compare phase to gaussian analysis. + z0 = n*dpix * np.tan(np.deg2rad(30)) + + z = np.linspace(0, 2*z0, n) + print(int((3+np.sqrt(3))/3*n)) + datacube = [] + for i in range(len(z)): + print(i) + A = gauss.fundamental_mode(z[i],n,m,np.max(x)) + # _f_ip_real = interpolate.interp2d(x, y, A.real, kind='linear') + # _f_ip_imag = interpolate.interp2d(x, y, A.imag, kind='linear') + # A_out_ip_real = _f_ip_real(_xi, _yi) + # A_out_ip_imag = _f_ip_imag(_xi, _yi) + # A = _E_out_ip_real + _E_out_ip_imag*1j + datacube.append(A[:,i]) + + plt.imshow(np.angle(np.asarray(datacube)));plt.colorbar();plt.show() + plt.imshow(np.abs(np.asarray(datacube)));plt.colorbar();plt.show() + set_trace() + ''' + + + + + + + + + + + + \ No newline at end of file diff --git a/scripts/planewave_rotation_test_2020_2.py b/scripts/planewave_rotation_test_2020_2.py new file mode 100644 index 0000000..8d8d818 --- /dev/null +++ b/scripts/planewave_rotation_test_2020_2.py @@ -0,0 +1,362 @@ +#!/usr/bin/env Python3 +import numpy as np +from pdb import set_trace +import matplotlib.pyplot as plt +from matplotlib import cm, colors +from cmath import sqrt as csqrt +from scipy import interpolate +from scipy.optimize import leastsq +from matplotlib.ticker import FormatStrFormatter +from matplotlib.colors import LogNorm +import pyfftw +# import glassCatalog +from tabulate import tabulate +import types +import time +from math import ceil +import sys +from mpl_toolkits import mplot3d +sys.path.insert(0, '../../') # where fourieroptics.py is located, so that it can be imported +import fourierOptics as fouOpt + + + +class FourierOptics: + def __init__(self, l): + self.l = l # Wavelength + + + def plane_wave_propagation_rotation(self, A, dpix=None, Rx=None, Ry=None, oversample=None, show=False, save=False, scaleInputArray=False, outputLimits=False): + ''' + Propagation through the Fresnel approximation, which is valid in many different situations. + A separate code for the Fraunhofer propagation is not striclty necessary, as the exponent + converges to 1 for large z values in any case. + + Input parameters: + A - [array] (Complex) Input field + dpix - [float] pixel scale at which the calculations are done + dz - [float] propagation distance + oversample - [bool] Apply padding on the outside of A, increasing the size of A + by the given factor. + show - [bool] Show the resulting amplitude and phase map after propagation + save - [bool] Save the figure + scaleInputArray - [tuple] Rescale the resolution of the input array using interpolation. + outputCoordinates - [list] list the corners of the desired output coordinate plane. + Order: [xmin, xmax, ymin, ymax] + + Output: + E_out - [array] (Complex) Output field distribution + x, y - [array] Output coordinates + ''' + + if scaleInputArray: + A = np.asarray(A) + _n, _m = A.shape + _nA, _mA = np.linspace(-1, 1, _n), np.linspace(-1, 1, _m) + A_ip_real = interpolate.interp2d(_nA, _mA, A.real, kind='cubic') + A_ip_imag = interpolate.interp2d(_nA, _mA, A.imag, kind='linear') + + _a, _b = scaleInputArray + _nx, _my = np.linspace(-1, 1, _a), np.linspace(-1, 1, _b) + A_out_ip_real = A_ip_real(_nx, _my) + A_out_ip_imag = A_ip_imag(_nx, _my) + A_out_ip = A_out_ip_real + A_out_ip_imag*1j + A = np.rint(A_out_ip_real) + + if oversample: + _Aos = np.zeros((np.shape(A)[0]*oversample, np.shape(A)[1]*oversample)) + _Aos[int(len(_Aos)/2-len(A)/2): int(len(_Aos)/2+len(A)/2), int(len(_Aos)/2-len(A)/2): int(len(_Aos)/2+len(A)/2)] = A + A = _Aos + + _n, _m = np.shape(A) + + # A at z=0 is then + _E_xy0 = pyfftw.empty_aligned((_n,_m), dtype='complex128') + _E_xy0[:] = np.fft.ifftshift(A) # use ifftshift to shift origin to start array + _IntensityNormFactor = np.sum(np.abs(_E_xy0)**2) # Factor to scale the intensity after the FFTs + _Az0 = pyfftw.empty_aligned((_n,_m), dtype='complex128') + _FFTW_Obj = pyfftw.FFTW( _E_xy0, _Az0, axes=(0,1), direction='FFTW_BACKWARD', threads=6, flags=('FFTW_ESTIMATE',)) + _FFTW_Obj() + + # E-field at z=dz is the fourier transform of A(z=0)*e^(-ikz) + _fx_l = np.fft.fftfreq(_n, d=dpix) + _fy_l = np.fft.fftfreq(_m, d=dpix) + _fx, _fy = np.meshgrid(_fx_l, _fy_l) + _fz = np.sqrt( (self.l**(-2) - _fx**2 - _fy**2).astype('complex128') ) + + # Rotation about the x and y axis from radians to degrees + Rx = np.deg2rad(Rx) + Ry = np.deg2rad(Ry) + + # Rotation matrices + _Rx = np.array([[1,0,0], + [0, np.cos(Rx), -np.sin(Rx)], + [0, np.sin(Rx), np.cos(Rx)]]) + _Ry = np.array([[np.cos(Ry), 0, np.sin(Ry)], + [0,1,0], + [-np.sin(Ry), 0, np.cos(Ry)]]) + _R = np.dot(_Rx,_Ry) + + # Rotation in k-space + _fxr = (_R[0,0]*_fx + _R[1,0]*_fy + _R[2,0]*_fz) + _fyr = (_R[0,1]*_fx + _R[1,1]*_fy + _R[2,1]*_fz) + _fzr = np.sqrt( (self.l**(-2) - _fxr**2 - _fyr**2).astype('complex128') ) + + # Carrier frequency + _u0 = _R[2,0]*1/self.l + _v0 = _R[2,1]*1/self.l + + # Recenter the new rotated frequencies at 0,0 + _fxr_c = (_fxr - _u0) + _fyr_c = (_fyr - _v0) + _fzr_c = np.sqrt( (self.l**(-2) - _fxr_c**2 - _fyr_c**2).astype('complex128') ) + + + # Calculation of the Jacobian + # _J = (_R[0,0]*_R[1,1] - _R[0,1]*_R[1,0]) + (_R[0,1]*_R[1,2] - _R[0,2]*_R[1,1])*_fxr_c/_fzr_c + (_R[0,2]*_R[1,0]-_R[0,0]*_R[1,2])*_fyr_c/_fzr_c + _J = (_R[0,0]*_R[1,1] - _R[0,1]*_R[1,0]) + (_R[1,0]*_R[2,0] - _R[2,0]*_R[1,1])*_fxr_c/_fzr_c + (_R[0,2]*_R[2,0]-_R[0,0]*_R[2,1])*_fyr_c/_fzr_c + + # Fourier transform of the input distribution, _A_in + _E_rot = pyfftw.empty_aligned((_n,_m), dtype='complex128') + _A_in = pyfftw.empty_aligned((_n,_m), dtype='complex128') + _A_in[:] = _Az0 * np.abs(_J) + _FFTW_Obj2 = pyfftw.FFTW( _A_in, _E_rot, axes=(0,1), direction='FFTW_BACKWARD', threads=6, flags=('FFTW_ESTIMATE',)) + _FFTW_Obj2() + + # Coordinates of E_rot + _y, _x = np.mgrid[-_n / 2 * dpix : _n / 2 * dpix : 1j*_n, -_m / 2 * dpix : _m / 2 * dpix : 1j * _m] + + # Multiplication with the phase contribution due to frequency shift, + # i.e. shift theorem. + _carr_freq = np.exp(1j * 2 * np.pi * (_u0 * _x + _v0 * _y)) + _E_rot = np.fft.fftshift(_E_rot) # use fftshift to move the origin from edge to center + _E_rot *= _carr_freq + _E_rot /= np.sqrt(_IntensityNormFactor) # Normalization back to initial energy. + + + if Rx or Ry: + # A rotation requires interpolation over the output + # Performing interpolation to get to the required final plane, done for real and imaginary part. + _xs = np.linspace(-_n/2*dpix, _n/2*dpix, _n) + _ys = np.linspace(-_m/2*dpix, _m/2*dpix, _m) + _zs = np.zeros(_xs.shape) + + _xr = _R[0,0]*_xs + _R[1,0]*_ys + _R[2,0]*_zs # Rotated original coordinates + _yr = _R[0,1]*_xs + _R[1,1]*_ys + _R[2,1]*_zs # Rotated original coordinates + _f_ip_real = interpolate.interp2d(_xs, _ys, _E_rot.real, kind='linear') + _f_ip_imag = interpolate.interp2d(_xs, _ys, _E_rot.imag, kind='linear') + # _f_ip_real = interpolate.interp2d(_xr, _yr, _E_rot.real, kind='linear') + # _f_ip_imag = interpolate.interp2d(_xr, _yr, _E_rot.imag, kind='linear') + _xi = np.linspace(_xr.min(), _xr.max(), _n) + _yi = np.linspace(_yr.min(), _yr.max(), _m) + # _xi = _xs + # _yi = _ys + + _E_out_ip_real = _f_ip_real(_xi, _yi) + _E_out_ip_imag = _f_ip_imag(_xi, _yi) + _E_rot = _E_out_ip_real + _E_out_ip_imag*1j + _y, _x = np.meshgrid(_yi, _xi) + + if show or save: + # Plot the amplitude and phase + _E_rot_extent = [np.min(_x), np.max(_x), np.min(_x), np.max(_x)] + fig = plt.figure(figsize=(8,3)) + ax = fig.add_subplot(121) + im = ax.imshow(np.abs(_E_rot), cmap='hot', extent=_E_rot_extent)#, vmax=1, vmin=0) + ax.set_xlabel('x [mm]') + ax.set_ylabel('y [mm]') + fig.colorbar(im, ax=ax, label='Normalized intensity') + ax2 = fig.add_subplot(122) + im2 = ax2.imshow(np.angle(_E_rot), cmap='hot', extent=_E_rot_extent) + ax2.set_xlabel('x [mm]') + ax2.set_ylabel('y [mm]') + ax.set_xlim([-0.4,0.4]) + ax.set_ylim([-0.4,0.4]) + ax2.set_xlim([-0.4,0.4]) + ax2.set_ylim([-0.4,0.4]) + fig.colorbar(im2, ax=ax2) + plt.tight_layout() + if save: + plt.savefig(save, bbox_inches='tight', transparent=True) + if show: + plt.show() + elif not show: + plt.close() + + return np.array([_E_rot, _x, _y]) + + + def phase_exponent(self, fx, fy, l): + k = 2*np.pi / l + kx = np.asarray(2*np.pi * fx) + ky = np.asarray(2*np.pi * fy) + K_squared = (kx**2 + ky**2) + kA = np.ones(K_squared.shape) * k + K_1 = np.sqrt((kA**2-K_squared).astype('complex128')) + kz = K_1.real + -1j*K_1.imag + return kz + + + def coordinate_rotation(self, Rx, Ry, n=5, m=5): + # Visualisation of the rotation matrix and projection. + # + # A circle C is plotted on the source frame (dark blue), + # the rotation of C onto the reference plane is shown in red + # and C is also projected from the source frame onto the + # reference frame in black. + # + _n = n + _m = m + Rx = np.deg2rad(Rx) + Ry = np.deg2rad(Ry) + dpix = 1/30 * self.l * 1000 + + _Rx = np.array([[1,0,0], + [0, np.cos(Rx), -np.sin(Rx)], + [0, np.sin(Rx), np.cos(Rx)]]) + _Ry = np.array([[np.cos(Ry), 0, np.sin(Ry)], + [0,1,0], + [-np.sin(Ry), 0, np.cos(Ry)]]) + _R = np.dot(_Rx,_Ry) + + # A grid in the source frame, coordinates (xs, ys, zs) + _ys, _xs = np.mgrid[-_n / 2 * dpix : _n / 2 * dpix : 1j*_n, -_m / 2 * dpix : _m / 2 * dpix : 1j * _m] + _zs = np.zeros(_xs.shape) + + # Rotated coordinates of the grid in the reference frame, coordinates (xr, yr, zr) + _xr = _R[0,0]*_xs + _R[1,0]*_ys + _R[2,0]*_zs # Rotated original coordinates + _yr = _R[0,1]*_xs + _R[1,1]*_ys + _R[2,1]*_zs # Rotated original coordinates + _zr = _R[0,2]*_xs + _R[1,2]*_ys + _R[2,2]*_zs + + # The circle C in the source frame. + theta = range(360) + _xcirc = 1.2 * np.cos(np.deg2rad(theta)) + _ycirc = 1.2 * np.sin(np.deg2rad(theta)) + _zcirc = np.zeros(_xcirc.shape) + + # Rotated coordinates + _xcirc_r = _R[0,0]*_xcirc + _R[1,0]*_ycirc + _R[2,0]*_zcirc # Rotated original coordinates + _ycirc_r = _R[0,1]*_xcirc + _R[1,1]*_ycirc + _R[2,1]*_zcirc # Rotated original coordinates + _zcirc_r = _R[0,2]*_xcirc + _R[1,2]*_ycirc + _R[2,2]*_zcirc + + # Projected coordinates + _xcirc_p = _xcirc + _ycirc_p = _ycirc + _zcirc_p = (-_R[2,0]*_xcirc_p - _R[2,1]*_ycirc_p) / _R[2,2] + + # Plot 3d Wireframe of source plane and reference plane + fig = plt.figure(figsize=(8,8)) + ax = plt.axes(projection='3d') + ax.plot_wireframe(_zs, _xs, _ys, label='source frame', color='blue', alpha=.3) + ax.plot_wireframe(_zr, _xr, _yr, label='reference frame', color='red', alpha=.3) + ax.plot3D(_zcirc, _xcirc, _ycirc, 'blue') + ax.plot3D(_zcirc_r, _xcirc_r, _ycirc_r, 'red', alpha=.2) + ax.plot3D(_zcirc_p, _xcirc_p, _ycirc_p, 'black') + ax.set_xlabel('z') + ax.set_ylabel('x') + ax.set_zlabel('y') + ax.set_xlim([-3,3]) + ax.set_ylim([-3,3]) + ax.set_zlim([-3,3]) + ax.grid(True, alpha=0.3) + ax.legend() + plt.tight_layout() + # plt.savefig('./coordinate_rotation_figures/coordinate_rotation_rx{:.0f}_ry{:.0f}.png'.format(np.degrees(Rx), np.degrees(Ry)), bbox_inches='tight') + # plt.close() + plt.show(block=False) + + + + +if __name__ == '__main__': + + # Pick a wavelength + l = 30e-3 + + # Initiate Fourier Optics + fou = FourierOptics(l=l) + fouOpt_norm = fouOpt.FourierOptics(l=l) + + # Pick sampling parameters + dpix = 1/15*l + n, m = 1001, 1001 + 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 = fouOpt.Profiles() + + # Use a Top Hat profile: + # A = prof.top_hat(4*l, n, m, dpix=dpix, oversample=1) + + # Or use a gaussian: + x, y = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m)) + gauss = fouOpt.GaussianBeam(0.06, l) + A = gauss.fundamental_mode(4*l,n,m,np.max(x)) + + # Definition of the rotation angles. + # Rotation about the x-axis is done first, then about the y-axis. + theta_x = 45 # Degrees + theta_y = 45 # Degrees + + # Plane wave propagation to 4*lambda + # Exyz,x4l,y4l = fouOpt_norm.plane_wave_propagation(A, dz=4*l, dpix=dpix, show=True) + + # Show the input distribution + _E_rot_extent = [np.min(x), np.max(x), np.min(y), np.max(y)] + fig = plt.figure(figsize=(8,3)) + ax = fig.add_subplot(121) + im = ax.imshow(np.abs(A), cmap='hot', extent=_E_rot_extent) + ax.set_xlabel('x [mm]') + ax.set_ylabel('y [mm]') + fig.colorbar(im, ax=ax, label='Normalized intensity') + ax2 = fig.add_subplot(122) + im2 = ax2.imshow(np.angle(A), cmap='hot', extent=_E_rot_extent) + ax2.set_xlabel('x [mm]') + ax2.set_ylabel('y [mm]') + fig.colorbar(im2, ax=ax2) + plt.tight_layout() + plt.show() + + + # Illustration of the rotation and projection for the given rotation angles. + fou.coordinate_rotation(theta_x,theta_y) + + # Plane wave propagation onto a rotated plane. + Erot,x,y = fou.plane_wave_propagation_rotation(A, dpix=dpix, Rx=theta_x, Ry=theta_y, show=True) + + + + + + ''' + # Compare phase to gaussian analysis. + z0 = n*dpix * np.tan(np.deg2rad(30)) + + z = np.linspace(0, 2*z0, n) + print(int((3+np.sqrt(3))/3*n)) + datacube = [] + for i in range(len(z)): + print(i) + A = gauss.fundamental_mode(z[i],n,m,np.max(x)) + # _f_ip_real = interpolate.interp2d(x, y, A.real, kind='linear') + # _f_ip_imag = interpolate.interp2d(x, y, A.imag, kind='linear') + # A_out_ip_real = _f_ip_real(_xi, _yi) + # A_out_ip_imag = _f_ip_imag(_xi, _yi) + # A = _E_out_ip_real + _E_out_ip_imag*1j + datacube.append(A[:,i]) + + plt.imshow(np.angle(np.asarray(datacube)));plt.colorbar();plt.show() + plt.imshow(np.abs(np.asarray(datacube)));plt.colorbar();plt.show() + ''' + + + + + + + + + + + + \ No newline at end of file diff --git a/scripts/psd_filtering_ELT_equiv.py b/scripts/psd_filtering_ELT_equiv.py new file mode 100644 index 0000000..df80364 --- /dev/null +++ b/scripts/psd_filtering_ELT_equiv.py @@ -0,0 +1,220 @@ +#!/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 sys +sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics') +from fourierOptics import * + + +''' + @ Author: Joost van den Born +''' + + +if __name__ == '__main__': + plt.rc('font', family='serif') + plt.rc('axes', labelsize='large') + + # Pick a wavelength + l = 1.35e-3 + + # Initiate Fourier Optics + fou = FourierOptics(l=l) + + # Pick sampling parameters + dpix = 160 + n, m = 1001, 1001 + oC = [-.5, .5, -.5, .5] + 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(38854/2, 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)) + + # Add PSD wavefront error to the lens + wfe = WavefrontError(n,m) + wfe.generate_height_map(5e-5, 2, 1) # 50 nm RMS for standard + h_data = wfe.data + # PSDSurface = surface(h_data) + + + # Setting up empty lists for fitting process + popt_airy_list, popt_cust_list = [], [] + pcov_airy_list, custom_var_list = [], [] + strehl_list = [] + Exyzlist, PSDlist, surfacelist, EEFlist = [], [], [], [] + + + # Create a lens with a focal length ~100 mm, with a clear aperture radius of 4 mm and thickness of 1 mm. + 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) + eef_og, r_og = fou.encircled_energy_function(np.abs(Exyz_og)**2, xp, yp) + + + + # Equidistant bins in LogSpace + flist = 10**np.linspace(np.log10(1), np.log10(x.shape[0]), 8) + + # Number of Zernike coefficients that may be fitted. + zernike = True + if zernike: + flist = [3, 6, 10, 15, 21, 28, 36, 45] + + print(flist) + set_trace() + + for fmax in flist: + + if zernike: + print('Sample {:02d}/{:02d}'.format(flist.index(fmax) + 1, len(flist) ) ) + wfe.fit_zernike_coefficients(h_data, fmax) + wfe.filteredData = wfe.zernike_cartesian(wfe.zernikeCoefficients, xp.real, yp.real) + if not zernike: + print('Sample {:02d}/{:02d}'.format(flist.tolist().index(fmax) + 1, len(flist) ) ) + if fmax == flist[0]: + # PSDSurface.set_lowpass_filter(fmax) + wfe.set_lowpass_filter(fmax) + else: + i = flist.tolist().index(fmax) - 1 + # PSDSurface.set_bandpass_filter(flist[i], fmax) + wfe.set_bandpass_filter(flist[i], fmax) + + set_trace() + newRMS = np.sqrt(np.mean(np.asarray(wfe.filteredData)**2))*1e6 + wfe.scale_rms(wfe.rms*1e6/newRMS, filteredDataFlag=False) + scal = wfe.scaledData + + set_trace() + # Adding the wavefront error + M1.add_wavefront_error(scal) + t_l = M1.phase_delay(x,y,l) + + # 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,xp,yp = fou.fresnel_propagation(U_in, x, y, dz=M1.f, show=False, outputCoordinates=oC) + eef, _ = fou.encircled_energy_function(np.abs(Exyz_psd)**2, xp, yp, originalPSF=np.abs(Exyz_og)**2) + + # Set up fitting object + psfFit_psd = PSFFitter(xp.real,yp.real,np.abs(Exyz_psd)**2) + + # Fit airy pattern, Gaussian and Moffat functions + airyFitted, popt_airy, pcov_airy = psfFit_psd.fit_airy_PSF(initialGuess=(0.03**2, 0.0, 0.0, 0.0, 0.03)) + custFitted, popt_cust, cust_var, cust_coupling = psfFit_psd.fit_custom_field(Exyz_og, Exyz_psd, return_coupling=True) + popt_airy_list.append(popt_airy) + popt_cust_list.append(popt_cust) + pcov_airy_list.append([ np.sqrt(pcov_airy[i][i]) for i in range(len(pcov_airy)) ]) + custom_var_list.append(np.sqrt(cust_var)) + + # Other properties + strehl_list.append(np.abs(cust_coupling)**2) + Exyzlist.append(Exyz_psd) + PSDlist.append(wfe.calculate_1D_PSD(filteredDataFlag=True)) + surfacelist.append(wfe.filteredData) + EEFlist.append(eef) + + + + + # Optimum fit parameters for all scaled surface errors + amplitude_airy, x0_airy, y0_airy, R_airy, offset_airy = np.asarray(popt_airy_list).T + x0_cust, y0_cust = np.asarray(popt_cust_list).T + + # 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 + 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_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 the centroid shift for all bandpasses + fig = plt.figure() + ax1 = fig.add_subplot(121) + ax2 = fig.add_subplot(122) + ax1.errorbar(flist, r_airy, yerr=r_err_airy, color='blue', marker='o', linestyle='--', label='Airy', capsize=3, markeredgewidth=1) + ax1.errorbar(flist, r_cust, yerr=r_err_cust, color='orange', marker='o', linestyle='--', label='Custom', capsize=3, markeredgewidth=1) + ax1.set_xscale('log') + ax1.set_xlabel('Filtered spatial frequency (lowpass)') + 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(block=False) + + + # Plotting the WFE, PSD and image planes + original_PSD_curve = wfe.calculate_1D_PSD(filteredDataFlag=False) + # fig2 = plt.figure() + for i in range(len(flist)): + fig2 = plt.figure(figsize=(12,4)) + # Plot of the surface profile + axi1 = fig2.add_subplot(1, 3, 1) + imi1 = axi1.imshow(surfacelist[i], cmap='jet', origin='lower') + fig2.colorbar(imi1, ax=axi1) + + # Plot of the PSD + axi2 = fig2.add_subplot(1, 3, 2) + axi2.loglog(original_PSD_curve, 'k-') + axi2.loglog(PSDlist[i], 'y-') + axi2.grid(which='major', alpha=0.3) + axi2.grid(which='minor', alpha=0.1) + + # Plot of the PSF + axi3 = fig2.add_subplot(1, 3, 3) + imi3 = axi3.imshow(np.abs(Exyzlist[i]), extent=oC, cmap='hot', origin='lower') + fig2.colorbar(imi3, ax=axi3) + axi3.set_xlim([-0.2,0.2]) + axi3.set_ylim([-0.2,0.2]) + + + plt.tight_layout() + # plt.savefig('fig_PSD_scaling_{}.png'.format(i+1), bbox_inches='tight') + # plt.close() + plt.show(block=False) + + + # Plotting the encircled energy for all iterations + fig3 = plt.figure() + ax = fig3.add_subplot(111) + ax.grid(alpha=0.3) + ax.set_xlabel('r') + ax.set_ylabel('Encircled energy fraction') + ax.plot(r_og, eef_og, 'r-', label='Diffraction limit') + for i in range(len(EEFlist)): + ax.plot(r_og, EEFlist[i], label='Filter #{}'.format(i+1)) + fig3.legend() + plt.show() + + + + + + + + + diff --git a/scripts/psd_filtering_ELT_zernikes.py b/scripts/psd_filtering_ELT_zernikes.py new file mode 100644 index 0000000..71555b5 --- /dev/null +++ b/scripts/psd_filtering_ELT_zernikes.py @@ -0,0 +1,233 @@ +#!/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 sys +sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics') +from fourierOptics import * + + +''' + @ Author: Joost van den Born +''' +def scale_zernike(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 + + +if __name__ == '__main__': + plt.rc('font', family='serif') + plt.rc('axes', labelsize='large') + + # Pick a wavelength + l = 1.35e-3 + + # Initiate Fourier Optics + fou = FourierOptics(l=l) + + # Pick sampling parameters + dpix = 160 + n, m = 1001, 1001 + oC = [-.5, .5, -.5, .5] + 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(36905/2, 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)) + + # Add PSD wavefront error to the lens + wfe = WavefrontError(n,m) + + # Generate some zernike coefficients + coeffic = np.abs(np.random.normal(0.5,0.3,28)).tolist()[::-1] + + yz,xz = np.mgrid[-1:1:wfe.n*1j, -1:1:wfe.m*1j] + _r = np.sqrt(xz**2 + yz**2) + _mask = _r > 0.95 + + rms_zern = 50e-5 # 50 nm + coeffic = scale_zernike(coeffic, rms_zern) + h_zern = wfe.zernike_cartesian(coeffic, xz, yz) + h_data = wfe.data + # print(coeffic) + + + # Setting up empty lists for fitting process + popt_airy_list, popt_cust_list = [], [] + pcov_airy_list, custom_var_list = [], [] + strehl_list = [] + Exyzlist, PSDlist, surfacelist, EEFlist = [], [], [], [] + + + # Create a lens with a focal length ~100 mm, with a clear aperture radius of 4 mm and thickness of 1 mm. + 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) + eef_og, r_og = fou.encircled_energy_function(np.abs(Exyz_og)**2, xp, yp) + + + # Number of Zernike coefficients that may be fitted. + zernike = True + if zernike: + flist = [1, 3, 6, 10, 15, 21, 28, 36, 45] + + + + for fmax in flist: + if zernike: + # SCALE THE COEFFICIENTS SUCH THAT THE RMS IS ALWAYS THE SAME??? + print('Sample {:02d}/{:02d}'.format(flist.index(fmax) + 1, len(flist) ) ) + yz, xz = np.mgrid[-1:1:wfe.n*1j, -1:1:wfe.m*1j] + wfe.fit_zernike_coefficients(h_data, xz, yz, modes=fmax) + wfe.filteredData = wfe.zernike_cartesian(wfe.zernikeCoefficients, xz, yz) + + if not zernike: + print('Sample {:02d}/{:02d}'.format(flist.tolist().index(fmax) + 1, len(flist) ) ) + if fmax == flist[0]: + # PSDSurface.set_lowpass_filter(fmax) + wfe.set_lowpass_filter(fmax) + else: + i = flist.tolist().index(fmax) - 1 + # PSDSurface.set_bandpass_filter(flist[i], fmax) + wfe.set_bandpass_filter(flist[i], fmax) + + newRMS = np.sqrt(np.mean(np.asarray(wfe.filteredData)**2))*1e6 + wfe.scale_rms(wfe.rms*1e6/newRMS, filteredDataFlag=False) + scal = wfe.scaledData + + + # Adding the wavefront error + M1.add_wavefront_error(wfe.filteredData) + t_l = M1.phase_delay(x,y,l) + + # 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,xp,yp = fou.fresnel_propagation(U_in, x, y, dz=M1.f, show=False, outputCoordinates=oC) + eef, _ = fou.encircled_energy_function(np.abs(Exyz_psd)**2, xp, yp, originalPSF=np.abs(Exyz_og)**2) + + # Set up fitting object + psfFit_psd = PSFFitter(xp.real,yp.real,np.abs(Exyz_psd)**2) + + # Fit airy pattern, Gaussian and Moffat functions + airyFitted, popt_airy, pcov_airy = psfFit_psd.fit_airy_PSF(initialGuess=(0.03**2, 0.0, 0.0, 0.0, 0.03)) + custFitted, popt_cust, cust_var, cust_coupling = psfFit_psd.fit_custom_field(Exyz_og, Exyz_psd, return_coupling=True) + popt_airy_list.append(popt_airy) + popt_cust_list.append(popt_cust) + pcov_airy_list.append([ np.sqrt(pcov_airy[i][i]) for i in range(len(pcov_airy)) ]) + custom_var_list.append(np.sqrt(cust_var)) + + # Other properties + strehl_list.append(np.abs(cust_coupling)**2) + Exyzlist.append(Exyz_psd) + PSDlist.append(wfe.calculate_1D_PSD(filteredDataFlag=True)) + surfacelist.append(wfe.filteredData) + EEFlist.append(eef) + + + + + # Optimum fit parameters for all scaled surface errors + amplitude_airy, x0_airy, y0_airy, R_airy, offset_airy = np.asarray(popt_airy_list).T + x0_cust, y0_cust = np.asarray(popt_cust_list).T + + # 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 + 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_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 the centroid shift for all bandpasses + fig = plt.figure() + ax1 = fig.add_subplot(121) + ax2 = fig.add_subplot(122) + ax1.errorbar(flist, r_airy, yerr=r_err_airy, color='blue', marker='o', linestyle='--', label='Airy', capsize=3, markeredgewidth=1) + ax1.errorbar(flist, r_cust, yerr=r_err_cust, color='orange', marker='o', linestyle='--', label='Custom', capsize=3, markeredgewidth=1) + ax1.set_xscale('log') + ax1.set_xlabel('Filtered spatial frequency (lowpass)') + 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(block=False) + + + # Plotting the WFE, PSD and image planes + original_PSD_curve = wfe.calculate_1D_PSD(filteredDataFlag=False) + # fig2 = plt.figure() + for i in range(len(flist)): + fig2 = plt.figure(figsize=(12,4)) + # Plot of the surface profile + axi1 = fig2.add_subplot(1, 3, 1) + imi1 = axi1.imshow(surfacelist[i], cmap='jet', origin='lower') + fig2.colorbar(imi1, ax=axi1) + + # Plot of the PSD + axi2 = fig2.add_subplot(1, 3, 2) + axi2.loglog(original_PSD_curve, 'k-') + axi2.loglog(PSDlist[i], 'y-') + axi2.grid(which='major', alpha=0.3) + axi2.grid(which='minor', alpha=0.1) + + # Plot of the PSF + axi3 = fig2.add_subplot(1, 3, 3) + imi3 = axi3.imshow(np.abs(Exyzlist[i]), extent=oC, cmap='hot', origin='lower') + fig2.colorbar(imi3, ax=axi3) + axi3.set_xlim([-0.2,0.2]) + axi3.set_ylim([-0.2,0.2]) + + + plt.tight_layout() + # plt.savefig('fig_PSD_scaling_{}.png'.format(i+1), bbox_inches='tight') + # plt.close() + plt.show(block=False) + + + # Plotting the encircled energy for all iterations + fig3 = plt.figure() + ax = fig3.add_subplot(111) + ax.grid(alpha=0.3) + ax.set_xlabel('r') + ax.set_ylabel('Encircled energy fraction') + ax.plot(r_og, eef_og, 'r-', label='Diffraction limit') + for i in range(len(EEFlist)): + ax.plot(r_og, EEFlist[i], label='Filter #{}'.format(i+1)) + fig3.legend() + plt.show() + + + + + + + + + diff --git a/scripts/psd_scaling_test.py b/scripts/psd_scaling_test.py new file mode 100644 index 0000000..cb1a8b3 --- /dev/null +++ b/scripts/psd_scaling_test.py @@ -0,0 +1,190 @@ +#!/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 sys +sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics') +from fourierOptics import * + + +''' + @ Author: Joost van den Born + @ Desc.: Importing a grid sag file and scaling it appropriately +''' + +class surface: + def __init__(self,data): + self.data = data + self.resolution = np.shape(self.data) + self.rms = np.sqrt(np.mean(np.asarray(self.data)**2)) + self.ptv = np.abs(np.max(self.data)) + np.abs(np.min(self.data)) + + def scale(self,factor): + # Scale the surface data such that the RMS gets multiplied by factor + _factor = factor ** 2 + self.PSD2Dmethod() + self.scaledPSD2D = self.PSD2D * _factor + # Fourier transform of the surface profile + _in = pyfftw.empty_aligned((n,m), dtype='complex128') + _in[:] = self.data + _out = pyfftw.empty_aligned((n,m), dtype='complex128') + _FFTW_Obj = pyfftw.FFTW( _in, _out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',)) + _FFTW_Obj() + _phase = np.angle(_out) + + _in = pyfftw.empty_aligned((n,m), dtype='complex128') + _in[:] = np.exp(1j*_phase) * np.sqrt(np.fft.fftshift(self.scaledPSD2D)) + _out = pyfftw.empty_aligned((n,m), dtype='complex128') + _FFTW_Obj = pyfftw.FFTW( _in, _out, axes=(0,1), direction='FFTW_BACKWARD', threads=6, flags=('FFTW_ESTIMATE',)) + _FFTW_Obj() + self.scaledData = _out.real + self.scaledrms = np.sqrt(np.mean(np.asarray(self.scaledData)**2)) + self.scaledptv = np.abs(np.max(self.scaledData)) + np.abs(np.min(self.scaledData)) + + def PSD2Dmethod(self): + ''' Returns the 2D PSD map ''' + _in = pyfftw.empty_aligned((n,m), dtype='complex128') + _in[:] = self.data + _out = pyfftw.empty_aligned((n,m), dtype='complex128') + _FFTW_Obj = pyfftw.FFTW( _in, _out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',)) + _FFTW_Obj() + self.D1 = np.fft.fftshift(_out) + self.PSD2D = np.abs(self.D1)**2 + return self.PSD2D + + +if __name__ == '__main__': + plt.rc('font', family='serif') + plt.rc('axes', labelsize='large') + + # Pick a wavelength + l = 30e-3 + + # Initiate Fourier Optics + fou = FourierOptics(l=l) + + # Pick sampling parameters + dpix = 1/50 + n, m = 1001, 1001 + 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)) + oC = [x.min(), x.max(), y.min(), y.max()] + + # Add PSD wavefront error to the lens + wfe = WavefrontError() + h_data = wfe.generate_height_map(1e-2, 2, 1, n, m) + + PSDSurface = surface(h_data) + print(PSDSurface.rms) + + + popt_airy_list, popt_gauss_list, popt_moffat_list = [], [], [] + pcov_airy_list, pcov_gauss_list, pcov_moffat_list = [], [], [] + + Exyzlist = [] + # Scaling the surface profile by a factor f + flist = [0.01, 0.1, 0.2, 0.5, 0.75, 1, 1.5, 2, 2.5] + for f in flist: + PSDSurface.scale(f) + scal = PSDSurface.scaledData + print(f, PSDSurface.scaledrms) + + + # 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) + L1.add_wavefront_error(scal) + t_l = L1.phase_delay(x,y,l, paraxial=False) + # print('Focal length of lens {} is {:.1f} mm.'.format(L1.name, L1.f)) + + # 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. + savename = None + Exyz_psd,_,_ = fou.fresnel_propagation(U_in, x, y, dz=L1.f, show=False, save=savename, outputCoordinates=oC) + + # Fit airy pattern + psfFit_psd = PSFFitter(x,y,np.abs(Exyz_psd)**2) + airyFitted, popt_airy, pcov_airy = psfFit_psd.fit_airy_PSF(initialGuess=(0.01, 0.0, 0.0, 0.0, 0.5)) + gaussianFitted, popt_gauss, pcov_gauss = psfFit_psd.fit_gaussian_PSF(initialGuess=(0.01, 0, 0, 0.07, 0.07, 0, 0)) + moffatFitted, popt_moffat, pcov_moffat = psfFit_psd.fit_moffat_PSF(initialGuess=(0, 0, 0, 0.01, 0.01, 0, 0, 1)) + + popt_airy_list.append(popt_airy) + popt_gauss_list.append(popt_gauss) + popt_moffat_list.append(popt_moffat) + + 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)) ]) + pcov_moffat_list.append([ np.sqrt(pcov_moffat[i][i]) for i in range(len(pcov_moffat)) ]) + Exyzlist.append(Exyz_psd) + + + # Optimum fit parameters for all scaled surface errors + 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 + + # 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 + + # 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)) + + set_trace() + + fig = plt.figure() + ax1 = fig.add_subplot(121) + ax2 = fig.add_subplot(122) + ax1.errorbar(flist, r_airy, yerr=r_err_airy, color='blue', marker='o', linestyle='--', label='Airy', capsize=3, markeredgewidth=1) + ax1.errorbar(flist, r_gauss, yerr=r_err_gauss, color='red', marker='o', linestyle='--', label='Gaussian', capsize=3, markeredgewidth=1) + ax1.errorbar(flist, r_moffat, yerr=r_err_moffat, color='green', marker='o', linestyle='--', label='Moffat', capsize=3, markeredgewidth=1) + ax1.set_xscale('log') + ax1.set_xlabel('Scaling factor') + ax1.set_ylabel('Distance') + ax1.legend() + ax1.grid(True, alpha=0.3) + + ax2.errorbar(x0_airy, y0_airy, xerr=x0_airy_err, yerr=y0_airy_err, color='k', marker='o', linestyle='--', label='Airy') + ax2.set_aspect('equal') + ax2.grid(True, alpha=0.3) + plt.tight_layout() + plt.show(block=False) + + fig2 = plt.figure() + for i in range(len(flist)): + axi = fig2.add_subplot(4,3,i+1) + imi = axi.imshow(np.abs(Exyzlist[i])**2, extent=oC, cmap='jet') + fig2.colorbar(imi, ax=axi) + axi.title.set_text('Scaling factor: {}'.format(flist[i])) + axi.set_xlim([-2.5,2.5]) + axi.set_ylim([-2.5,2.5]) + plt.show() + + + + + + + + diff --git a/scripts/psd_scaling_test_ELT_equiv.py b/scripts/psd_scaling_test_ELT_equiv.py new file mode 100644 index 0000000..48e32a8 --- /dev/null +++ b/scripts/psd_scaling_test_ELT_equiv.py @@ -0,0 +1,406 @@ +#!/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 sys +sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics') +from fourierOptics import * + + +''' + @ Author: Joost van den Born + @ Desc.: Importing a grid sag file and scaling it appropriately +''' + +class surface: + def __init__(self,data): + self.data = data + self.resolution = np.shape(self.data) + self.rms = np.sqrt(np.mean(np.asarray(self.data)**2)) + self.ptv = np.abs(np.max(self.data)) + np.abs(np.min(self.data)) + + def scale_rms(self,factor): + # Scale the surface data such that the RMS gets multiplied by factor + _factor = factor ** 2 + self.calculate_2D_PSD() + self.scaledPSD2D = self.PSD2D * _factor + # Fourier transform of the surface profile + _in = pyfftw.empty_aligned((n,m), dtype='complex128') + _in[:] = self.data + _out = pyfftw.empty_aligned((n,m), dtype='complex128') + _FFTW_Obj = pyfftw.FFTW( _in, _out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',)) + _FFTW_Obj() + _phase = np.angle(_out) + + _in = pyfftw.empty_aligned((n,m), dtype='complex128') + _in[:] = np.exp(1j*_phase) * np.sqrt(np.fft.fftshift(self.scaledPSD2D)) + _out = pyfftw.empty_aligned((n,m), dtype='complex128') + _FFTW_Obj = pyfftw.FFTW( _in, _out, axes=(0,1), direction='FFTW_BACKWARD', threads=6, flags=('FFTW_ESTIMATE',)) + _FFTW_Obj() + self.scaledData = _out.real + self.scaledrms = np.sqrt(np.mean(np.asarray(self.scaledData)**2)) + self.scaledptv = np.abs(np.max(self.scaledData)) + np.abs(np.min(self.scaledData)) + + def calculate_2D_PSD(self): + ''' Returns the 2D PSD map ''' + _in = pyfftw.empty_aligned((n,m), dtype='complex128') + _in[:] = self.data + _out = pyfftw.empty_aligned((n,m), dtype='complex128') + _FFTW_Obj = pyfftw.FFTW( _in, _out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',)) + _FFTW_Obj() + self.D1 = np.fft.fftshift(_out) + self.PSD2D = np.abs(self.D1)**2 + return self.PSD2D + + def calculate_1D_PSD(self, filteredDataFlag=False): + ''' Returns the azimuthally averaged PSD map ''' + try: + if not filteredDataFlag: + self.PSD1D = self.azimuthal_average() + else: + self.filteredPSD1D = self.azimuthal_average(filteredDataFlag=filteredDataFlag) + return self.filteredPSD1D + except AttributeError: + self.calculate_2D_PSD() + self.PSD1D = self.azimuthal_average() + + return self.PSD1D + + + def set_lowpass_filter(self, r, dpix=1): + ''' Returns a map where all frequencies higher than r have been filtered out ''' + self.calculate_2D_PSD() + # Masking out the parts that fall outside the lens clear aperture + _dist_from_center = np.sqrt(x**2 + y**2) * dpix + _mask = (_dist_from_center <= (r)).astype(int) + _maskedPSD2D = self.PSD2D.copy() + _maskedPSD2D *= _mask + _phase = np.angle(np.fft.fft2(self.data)) + self.filteredPSD2D = _maskedPSD2D + self.filteredData = np.fft.ifft2(np.exp(1j*_phase) * np.sqrt(np.fft.fftshift(self.filteredPSD2D))).real + + def set_highpass_filter(self, r, dpix=1): + '''Returns a map where all spatial frequencies lower than r have been filtered out''' + self.calculate_2D_PSD() + # Masking out the parts that fall outside the lens clear aperture + x,y = np.ogrid[-self.resolution[0]/2:self.resolution[0]/2, -self.resolution[1]/2:self.resolution[1]/2] + _dist_from_center = np.sqrt(x**2 + y**2) * dpix + _mask = (_dist_from_center <= (r)).astype(int) + _maskedPSD2D = self.PSD2D.copy() + _maskedPSD2D *= (1-_mask) + _phase = np.angle(np.fft.fft2(self.data)) + self.filteredPSD2D = _maskedPSD2D + self.filteredData = np.fft.ifft2(np.exp(1j*_phase) * np.sqrt(np.fft.fftshift(self.filteredPSD2D))).real + + def set_bandpass_filter(self, rmin, rmax, dpix=1): + '''Returns a map where all spatial frequencies lower than r2 and higher than r1 have been filtered out''' + assert rmax > rmin, 'rmax < rmin, not physical!' + self.calculate_2D_PSD() + # Masking out the parts that fall outside the lens clear aperture + x,y = np.ogrid[-self.resolution[0]/2:self.resolution[0]/2, -self.resolution[1]/2:self.resolution[1]/2] + _dist_from_center = np.sqrt(x**2 + y**2) * dpix + _mask_1 = (_dist_from_center <= (rmax)).astype(int) + _mask_2 = (_dist_from_center >= (rmin)).astype(int) + _maskedPSD2D = self.PSD2D.copy() + _maskedPSD2D *= (_mask_1 & _mask_2) + _phase = np.angle(np.fft.fft2(self.data)) + self.filteredPSD2D = _maskedPSD2D + self.filteredData = np.fft.ifft2(np.exp(1j*_phase) * np.sqrt(np.fft.fftshift(self.filteredPSD2D))).real + + def azimuthal_average(self, center=None, filteredDataFlag=False): + """ + Calculate the azimuthally averaged radial profile. + + image - The 2D image + center - The [x,y] pixel coordinates used as the center. The default is + None, which then uses the center of the image (including + fracitonal pixels). + + """ + + # Calculate the indices from the image + _y, _x = np.indices(self.resolution) + if not center: + center = np.array([(_x.max()-_x.min())/2.0, (_y.max()-_y.min())/2.0]) + + _r = np.hypot(_x - center[0], _y - center[1]) + + # Get sorted radii + _ind = np.argsort(_r.flat) + _r_sorted = _r.flat[_ind] + if not filteredDataFlag: + _i_sorted = self.PSD2D.flat[_ind] + else: + _i_sorted = self.filteredPSD2D.flat[_ind] + + # Get the integer part of the radii (bin size = 1) + _r_int = _r_sorted.astype(int) + + # Find all pixels that fall within each radial bin. + _deltar = _r_int[1:] - _r_int[:-1] # Assumes all radii represented + _rind = np.where(_deltar)[0] # location of changed radius + _nr = _rind[1:] - _rind[:-1] # number of radius bin + + # Cumulative sum to figure out sums for each radius bin + _csim = np.cumsum(_i_sorted, dtype=float) + _tbin = _csim[_rind[1:]] - _csim[_rind[:-1]] + _radial_prof = _tbin / _nr + + return _radial_prof + + + +if __name__ == '__main__': + plt.rc('font', family='serif') + plt.rc('axes', labelsize='large') + + # Pick a wavelength + l = 1.35e-3 + + # Initiate Fourier Optics + fou = FourierOptics(l=l) + + # Pick sampling parameters + dpix = 160 + n, m = 1001, 1001 + oC = [-.5, .5, -.5, .5] + 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(38854/2, 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)) + # oC = [x.min(), x.max(), y.min(), y.max()] + + # Add PSD wavefront error to the lens + wfe = WavefrontError() + h_data = wfe.generate_height_map(5e-5, 2, 1, n, m) # 50 nm RMS for standard + + PSDSurface = surface(h_data) + # print(PSDSurface.rms) + # PSDSurface.set_bandpass_filter(10, 20) + # PSDSurface.calculate_1D_PSD(filteredDataFlag=False) + # PSDSurface.calculate_1D_PSD(filteredDataFlag=True) + + # fig = plt.figure() + # ax1 = fig.add_subplot(121) + # ax2 = fig.add_subplot(122) + # im1 = ax1.imshow(PSDSurface.filteredData, origin='lower', cmap='jet') + # ax2.loglog(PSDSurface.PSD1D) + # ax2.loglog(PSDSurface.filteredPSD1D) + # ax2.grid(alpha=0.3, which='major') + # ax2.grid(alpha=0.1, which='minor') + # fig.colorbar(im1, ax=ax1) + # plt.tight_layout() + # plt.show() + + 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 = [] + Exyzlist = [] + # Scaling the surface profile by a factor f + + + + + ''' + + + M1 = ParabolicMirror(dx=0, dy=0, name='ELT-M1', CA=38854/2, f=679945) + t_l = M1.phase_delay(x,y,l) + + # The input distribution gets a curved wavefront after the lens + U_in = A * t_l + + Exyz_og,xp,yp = fou.fresnel_propagation(U_in, x, y, dz=M1.f, show=False, outputCoordinates=oC) + + # Preparing some figures + fig1 = plt.figure() + ax = fig1.add_subplot(111) + im = ax.imshow(np.abs(Exyz_og)**2, cmap=plt.cm.jet, origin='bottom', + extent=oC) + fig1.colorbar(im) + + fig2 = plt.figure() + ax21 = fig2.add_subplot(121) + ax21.semilogy((np.abs(Exyz_og)**2)[500], 'k-', label='Original') + ax21.grid(True, alpha=0.3) + + + ax22 = fig2.add_subplot(122) + ax22.grid(True, alpha=0.3) + + + fig3 = plt.figure() + ax31 = fig3.add_subplot(131) + ax32 = fig3.add_subplot(132) + ax33 = fig3.add_subplot(133) + + # Start of fitting routine + psfFit = PSFFitter(xp.real,yp.real,np.abs(Exyz_og)**2) + + gaussianFit = moffatFit = airyFit = True + + # Fit a Gaussian profile + if gaussianFit: + gaussianFitted, popt_gauss, pcov_gauss = psfFit.fit_gaussian_PSF(initialGuess=(0.03**2, 0, 0, 0.02, 0.02, 0, 0)) + A_gauss, x0_gauss, y0_gauss, sigma_x_gauss, sigma_y_gauss, theta_gauss, offset_gauss = popt_gauss + ax.contour(xp.real, yp.real, gaussianFitted, 8, colors='w', linestyles='solid') + ax.plot(x0_gauss,y0_gauss, 'wx') + ax21.semilogy(gaussianFitted[500], 'r-', label='Gaussian fit') + ax22.semilogy((np.abs(Exyz_og)**2)[500] - gaussianFitted[500], 'r-', label='Gaussian differential') + im31 = ax31.imshow(np.abs(Exyz_og)**2-gaussianFitted, cmap=plt.cm.jet, origin='bottom', extent=oC) + fig3.colorbar(im31, ax=ax31) + + if moffatFit: + moffatFitted, popt_moffat, pcov_moffat = psfFit.fit_moffat_PSF(initialGuess=(0.03**2, 0, 0, 0.02, 0.02, 0, 0, 4)) + A_moffat, x0_moffat, y0_moffat, sigma_x_moffat, sigma_y_moffat, theta_moffat, offset_moffat, beta_moffat = popt_moffat + ax.contour(xp.real, yp.real, moffatFitted, 8, colors='y', linestyles='dashed') + ax.plot(x0_moffat,y0_moffat, 'yx') + ax21.semilogy(moffatFitted[500], 'g-', label='Moffat fit') + ax22.semilogy((np.abs(Exyz_og)**2)[500] - moffatFitted[500], 'g-', label='Moffat differential') + im32 = ax32.imshow(np.abs(Exyz_og)**2-moffatFitted, cmap=plt.cm.jet, origin='bottom', extent=oC) + fig3.colorbar(im32, ax=ax32) + + if airyFit: + airyFitted, popt_airy, pcov_airy = psfFit.fit_airy_PSF(initialGuess=(0.03**2, 0.0, 0.0, 0.0, 0.03)) + amplitude_airy, x0_airy, y0_airy, R_airy, offset_airy = popt_airy + ax.contour(xp.real, yp.real, airyFitted, 8, colors='c', linestyles='dashdot') + ax.plot(x0_airy,y0_airy, 'cx') + ax21.semilogy(airyFitted[500], 'b-', label='Airy fit') + ax22.semilogy((np.abs(Exyz_og)**2)[500] - airyFitted[500], 'b-', label='Airy differential') + im33 = ax33.imshow(np.abs(Exyz_og)**2-airyFitted, cmap=plt.cm.jet, origin='bottom', extent=oC) + fig3.colorbar(im33, ax=ax33) + + ax21.legend() + ax22.legend() + plt.tight_layout() + plt.show() + + + + + ''' + + + # Create a lens with a focal length ~100 mm, with a clear aperture radius of 4 mm and thickness of 1 mm. + 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 + Exyz_og,xp,yp = fou.fresnel_propagation(U_in, x, y, dz=M1.f, show=False, outputCoordinates=oC) + + eef2, r2 = fou.encircled_energy_function(np.abs(Exyz_og)**2, xp, yp) + plt.plot(r2, eef2); plt.grid();plt.show() + set_trace() + RMSlist = [] + flist = [0.01, 0.1, 0.2, 0.5, 0.75, 1, 1.5, 2, 3, 4, 5, 6] + for f in flist: + PSDSurface.scale_rms(f) + scal = PSDSurface.scaledData + print(f, '{:.1f} nm'.format(PSDSurface.scaledrms*1e6)) + RMSlist.append(PSDSurface.scaledrms*1e6) + M1.add_wavefront_error(scal) + t_l = M1.phase_delay(x,y,l) + + # 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,xp,yp = fou.fresnel_propagation(U_in, x, y, dz=M1.f, show=False, outputCoordinates=oC) + + # Set up fitting object + psfFit_psd = PSFFitter(xp.real,yp.real,np.abs(Exyz_psd)**2) + + # Fit airy pattern, Gaussian and Moffat functions + airyFitted, popt_airy, pcov_airy = psfFit_psd.fit_airy_PSF(initialGuess=(0.03**2, 0.0, 0.0, 0.0, 0.03)) + gaussianFitted, popt_gauss, pcov_gauss = psfFit_psd.fit_gaussian_PSF(initialGuess=(0.04, 0, 0, 0.01, 0.01, 0, 0)) + # moffatFitted, popt_moffat, pcov_moffat = psfFit_psd.fit_moffat_PSF(initialGuess=(0.04, 0, 0, 0.01, 0.01, 0, 0, 4)) + custFitted, popt_cust, cust_var, cust_coupling = psfFit_psd.fit_custom_field(Exyz_og, Exyz_psd, return_coupling=True) + + + popt_airy_list.append(popt_airy) + popt_gauss_list.append(popt_gauss) + # popt_moffat_list.append(popt_moffat) + 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)) ]) + # pcov_moffat_list.append([ np.sqrt(pcov_moffat[i][i]) for i in range(len(pcov_moffat)) ]) + custom_var_list.append(np.sqrt(cust_var)) + strehl_list.append(np.abs(cust_coupling)**2) + Exyzlist.append(Exyz_psd) + + + + # Optimum fit parameters for all scaled surface errors + 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)) + + 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(block=False) + + fig2 = plt.figure() + for i in range(len(flist)): + axi = fig2.add_subplot(4,3,i+1) + imi = axi.imshow(np.abs(Exyzlist[i])**2, extent=oC, cmap='jet') + # fig2.colorbar(imi, ax=axi) + axi.title.set_text('RMS: {:.1f} nm (S={:.2f})'.format(RMSlist[i], strehl_list[i])) + axi.set_xlim([-0.2,0.2]) + axi.set_ylim([-0.2,0.2]) + plt.tight_layout() + plt.show() + + + + + + + + diff --git a/scripts/simple_lens_system.py b/scripts/simple_lens_system.py new file mode 100644 index 0000000..ccee932 --- /dev/null +++ b/scripts/simple_lens_system.py @@ -0,0 +1,232 @@ +#!/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))) + + + diff --git a/scripts/verify_strehl.py b/scripts/verify_strehl.py new file mode 100644 index 0000000..9183fe7 --- /dev/null +++ b/scripts/verify_strehl.py @@ -0,0 +1,226 @@ +#!/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() +