fourierOptics/scripts/SPIE/wavefront_error_frequency_analysis.py

import numpy as np
import os
import sys
sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics')
from fourierOptics import *

'''
	Author: 		Joost van den Born, PhD student at University of Groningen
	Contact: 		born@astro.rug.nl
	Description: 	Script to analyse the shift of the PSF centroid as a 
					function of the RMS value of the WFE.
					Script made for the SPIE manuscript.
'''


if __name__ == '__main__':
	plt.rc('font', family='serif')
	plt.rc('axes', labelsize='large')

	# Pick a wavelength
	l = 1.49e-3
	input_field_location = './data/ADC_input_field_1001x1001_1.49_onaxis_circular.pkl'

	# Initializing MICADO paraxial model
	mcd = ParaxialMICADO(n=1001, m=1001, dpix=160, entrance_pupil='open')

	# Propagation of the unperturbed system.
	Exyz_og, xp, yp = mcd.propagate(l=l,surface=-1, zoom=False, show=False, 
									print_progress=False, hx=0, hy=0,
									load_up_to_adc=input_field_location)

	# Generate the wavefront error that will be used.
	wfe = WavefrontError(mcd.n,mcd.m)
	wfe.generate_height_map(5e-5, 2, 1) # 50 nm RMS for standard, power slope 2, cut-off 1
	h_data = wfe.data

	# Plotting the WFE, PSD and image planes
	oC = [np.min(xp), np.max(xp), np.min(yp), np.max(yp)]
	
	# Mask out areas where the ADC prism isn't
	_, xadc, yadc = mcd.propagate(l=l,surface=3, zoom=False, show=False, 
								print_progress=True, hx=0, hy=0)
	mask = np.sqrt(xadc**2 + yadc**2) > 62.5

	# Setting up empty lists for fitting process
	popt_airy_list, popt_cust_list = [], []
	pcov_airy_list, custom_var_list = [], []
	Exyzlist, PSDlist, surfacelist = [], [], []
	central_frequencies = []
	
	# Bins in LogSpace
	x = mcd.x
	flist = np.asarray([0.5, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 20, 25,
			 30, 35, 40, 45, 50, 60, 70, 80, 90, 100, 120, 140, 160, 180, 200,
			 220, 250, 280, 310, 340, 370, 400, 450, 500])
	plt.imshow(wfe.data*1e6, cmap='jet');plt.colorbar();plt.show()


	###############################
	### Lowpass filter analysis ###
	###############################

	# Do the analysis for each frequency in flist
	for fmax in flist:
		print('Lowpass filter - Freq.: {} - Sample {:02d}/{:02d}'.format(fmax, flist.tolist().index(fmax) + 1, len(flist) ) )
		if fmax == flist[0]:
			wfe.set_lowpass_filter(fmax)
		else:
			i = flist.tolist().index(fmax) - 1
			wfe.set_lowpass_filter(flist[i])

		scal = wfe.filteredData

		# Adding the wavefront error
		mcd.ELT_ADC.p1.add_wavefront_error(scal)

		# Performing the propagation using Fresnel approximation to the focal plane.
		Exyz_psd,xp,yp = mcd.propagate(l=l,surface=-1, zoom=False, show=False, 
											print_progress=False, hx=0, hy=0,
											load_up_to_adc=input_field_location)

		# Set up fitting object
		psfFit_psd = PSFFitter(xp.real,yp.real,np.abs(Exyz_psd)**2)

		# Fit airy pattern and coupling integral
		amplitude_guess = np.max(np.abs(Exyz_psd))**2
		airyFitted, popt_airy, pcov_airy = psfFit_psd.fit_airy_PSF(initialGuess=(amplitude_guess, 0.0, 0.0, 0.0, 0.03))
		popt_airy_list.append(popt_airy)
		pcov_airy_list.append([ np.sqrt(pcov_airy[i][i]) for i in range(len(pcov_airy)) ])

		# Other properties
		original_PSD_curve = wfe.PSD2D[int((mcd.n-1)/2),int((mcd.n-1)/2):]
		wfe.calculate_1D_PSD(filteredDataFlag=True)
		PSDlist.append(wfe.filteredPSD2D[int((mcd.n-1)/2),int((mcd.n-1)/2):])
		surfacelist.append(wfe.filteredData)

		# Plot of the surface profile
		fig2 = plt.figure(figsize=(12,3))
		axi1 = fig2.add_subplot(1, 3, 1)
		wavefrontmap = wfe.filteredData.copy()*1e6
		wavefrontmap[mask] = np.nan
		imi1 = axi1.imshow(wavefrontmap, cmap='jet', origin='lower', extent=[xadc.min(), xadc.max(), yadc.min(), yadc.max()])
		axi1.set_xlabel('x (mm)')
		axi1.set_ylabel('y (mm)')
		axi1.set_xlim([-65,65])
		axi1.set_ylim([-65,65])
		fig2.colorbar(imi1, ax=axi1, label='Surface height (nm)')

		# Plot of the PSD
		axi2 = fig2.add_subplot(1, 3, 2)
		axi2.loglog(original_PSD_curve, 'k-')
		axi2.loglog(PSDlist[-1], 'y-')
		axi2.grid(which='major', alpha=0.3)
		axi2.grid(which='minor', alpha=0.1)
		axi2.set_xlabel('Normalized Radial Frequency ($\\rho/\\rho_c$)')
		axi2.set_ylabel('PSD (mm$^4$)')

		# Plot of the PSF
		axi3 = fig2.add_subplot(1, 3, 3)
		differential_amplitude = np.abs(Exyz_og) - np.abs(Exyz_psd)
		if (differential_amplitude == 0).all():
			norm_par = None
		else:
			norm_par = LogNorm()
		imi3 = axi3.imshow(differential_amplitude, extent=oC, cmap='gray', origin='lower', norm=norm_par)
		fig2.colorbar(imi3, ax=axi3, label='Differential amplitude')
		axi3.set_xlim([-1,1])
		axi3.set_ylim([-1,1])
		axi3.set_xlabel('x (mm)')
		axi3.set_ylabel('y (mm)')
		plt.tight_layout()
		plt.savefig('./figures/wfe_frequency_analysis/lowpass_spatial_freq_{:.0f}.png'.format(flist.tolist().index(fmax) + 1), bbox_inches='tight', transparent=True)
		plt.close()


	# Optimum fit parameters for all scaled surface errors
	amplitude_airy, x0_airy, y0_airy, R_airy, offset_airy = np.asarray(popt_airy_list).T
	amplitude_airy_err, x0_airy_err, y0_airy_err, R_airy_err, offset_airy_err = np.asarray(pcov_airy_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_airy_mas = r_airy * mcd.platescale * 1000
	r_err_airy_mas = r_err_airy * mcd.platescale * 1000
	
	save_variable('./data/lowpass_frequency_filter_airy_distance.pkl', (r_airy, r_err_airy))
	save_variable('./data/lowpass_frequency_filter_airy_locations.pkl', (x0_airy, y0_airy, x0_airy_err, y0_airy_err))
	print('Platescale in mas/mm: ', mcd.platescale * 1000)



	################################
	### Bandpass filter analysis ###
	################################

	# Setting up empty lists for fitting process
	popt_airy_list, popt_cust_list = [], []
	pcov_airy_list, custom_var_list = [], []
	Exyzlist, PSDlist, surfacelist = [], [], []
	central_frequencies = []

	# Do the analysis for each frequency in flist
	for fmax in flist:
		print('Bandpass - Freq.: {} - Sample {:02d}/{:02d}'.format(fmax, flist.tolist().index(fmax) + 1, len(flist) ) )
		if fmax == flist[0]:
			# Use a lowpass filter only for the very first element in flist
			wfe.set_lowpass_filter(fmax)
			central_frequencies.append(fmax)
		else:
			# Set a bandpass between the current and previous frequency element
			i = flist.tolist().index(fmax) - 1
			wfe.set_bandpass_filter(flist[i], fmax)
			central_frequencies.append(np.mean([flist[i],fmax]))


		# Adding the wavefront error
		scal = wfe.filteredData
		mcd.ELT_ADC.p1.add_wavefront_error(scal)

		# Performing the propagation using Fresnel approximation to the focal plane.
		Exyz_psd,xp,yp = mcd.propagate(l=l,surface=-1, zoom=False, show=False, 
											print_progress=False, hx=0, hy=0,
											load_up_to_adc=input_field_location)

		# Set up fitting object
		psfFit_psd = PSFFitter(xp.real,yp.real,np.abs(Exyz_psd)**2)

		# Fit airy pattern and coupling integral
		amplitude_guess = np.max(np.abs(Exyz_psd))**2
		airyFitted, popt_airy, pcov_airy = psfFit_psd.fit_airy_PSF(initialGuess=(amplitude_guess, 0.0, 0.0, 0.0, 0.03))
		popt_airy_list.append(popt_airy)
		pcov_airy_list.append([ np.sqrt(pcov_airy[i][i]) for i in range(len(pcov_airy)) ])

		# Other properties
		original_PSD_curve = wfe.PSD2D[int((mcd.n-1)/2),int((mcd.n-1)/2):]
		wfe.calculate_1D_PSD(filteredDataFlag=True)
		PSDlist.append(wfe.filteredPSD2D[int((mcd.n-1)/2),int((mcd.n-1)/2):])
		surfacelist.append(wfe.filteredData)

		# Plot of the surface profile
		fig2 = plt.figure(figsize=(12,3))
		axi1 = fig2.add_subplot(1, 3, 1)
		wavefrontmap = wfe.filteredData.copy()*1e6
		wavefrontmap[mask] = np.nan
		imi1 = axi1.imshow(wavefrontmap, cmap='jet', origin='lower', extent=[xadc.min(), xadc.max(), yadc.min(), yadc.max()])
		axi1.set_xlabel('x (mm)')
		axi1.set_ylabel('y (mm)')
		axi1.set_xlim([-65,65])
		axi1.set_ylim([-65,65])
		fig2.colorbar(imi1, ax=axi1, label='Surface height (nm)')

		# Plot of the PSD
		axi2 = fig2.add_subplot(1, 3, 2)
		axi2.loglog(original_PSD_curve, 'k-')
		axi2.loglog(PSDlist[-1], 'y-')
		axi2.grid(which='major', alpha=0.3)
		axi2.grid(which='minor', alpha=0.1)
		axi2.set_xlabel('Normalized Radial Frequency ($\\rho/\\rho_c$)')
		axi2.set_ylabel('PSD (mm$^4$)')

		# Plot of the PSF
		axi3 = fig2.add_subplot(1, 3, 3)
		differential_amplitude = np.abs(Exyz_og) - np.abs(Exyz_psd)
		if (differential_amplitude == 0).all():
			norm_par = None
		else:
			norm_par = LogNorm()
		imi3 = axi3.imshow(differential_amplitude, extent=oC, cmap='gray', origin='lower', norm=norm_par)
		fig2.colorbar(imi3, ax=axi3, label='Differential amplitude')
		axi3.set_xlim([-1,1])
		axi3.set_ylim([-1,1])
		axi3.set_xlabel('x (mm)')
		axi3.set_ylabel('y (mm)')
		plt.tight_layout()
		plt.savefig('./figures/wfe_frequency_analysis/bandpass_spatial_freq_{:.0f}.png'.format(flist.tolist().index(fmax) + 1), bbox_inches='tight', transparent=True)
		plt.close()

	# Optimum fit parameters for all scaled surface errors
	amplitude_airy, x0_airy, y0_airy, R_airy, offset_airy = np.asarray(popt_airy_list).T
	amplitude_airy_err, x0_airy_err, y0_airy_err, R_airy_err, offset_airy_err = np.asarray(pcov_airy_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_airy_mas = r_airy * mcd.platescale * 1000
	r_err_airy_mas = r_err_airy * mcd.platescale * 1000
	
	save_variable('./data/bandpass_frequency_filter_airy_distance.pkl', (r_airy, r_err_airy))
	save_variable('./data/bandpass_frequency_filter_airy_locations.pkl', (x0_airy, y0_airy, x0_airy_err, y0_airy_err))
	save_variable('./data/bandpass_frequency_filter_central_frequencies.pkl', central_frequencies)
	print('Platescale in mas/mm: ', mcd.platescale * 1000)