fourierOptics/scripts/SPIE/geometric_distortions_incl_rotations.py

import numpy as np
import sys
sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Documents/Published/FourierOptics GitHub release/')
from fourierOptics import *

'''
	Author: 		Joost van den Born, PhD student at University of Groningen
	Contact: 		born@astro.rug.nl
	Description: 	Script to generate the geometric distortions due to WFEs
					off the ADC prisms. The distortions are produced by measuring
					the offset of the PSF centroid with and without WFE as a
					function of field position.
	 				Script made for the SPIE manuscript.
'''

# Rotation of ADC angles - [0.125, 0.25, 0.5, 0.75, 1, 2, 5, 10, 30, 45, 60, 90, 180]


def plot_distortions():
	# Coordinates without WFEs
	xe = load_variable('./data/ADC rotation/GD_obscured_PSF_xcentroid_expected_zern_psd.pkl')
	ye = load_variable('./data/ADC rotation/GD_obscured_PSF_ycentroid_expected_zern_psd.pkl')

	# Coordinates with WFEs
	xf = load_variable('./data/ADC rotation/GD_obscured_PSF_xcentroid_found_with_wfe_airy_zern_psd.pkl')
	yf = load_variable('./data/ADC rotation/GD_obscured_PSF_ycentroid_found_with_wfe_airy_zern_psd.pkl')

	mcd = ParaxialMICADO(n=3001, m=3001, dpix=50, l=1.49e-3)

	n = xe.shape[0]
	X = xe.flatten()
	Y = ye.flatten()
	dX = xf - xe
	dY = yf - ye
	dx = dX.flatten()
	dy = dY.flatten()

	# Magnitude of the distortions
	R = np.sqrt(dx**2 + dy**2) * mcd.platescale * 1e6

	# Dummy figure, used for the colorbar colors in the quiver plot
	im = plt.imshow(np.reshape(R, (n,n)))
	plt.close()

	fig = plt.figure()
	ax = fig.add_subplot()
	for i in range(len(xe)):
	    ax.plot(xe[:,i], ye[:,i], 'k-', alpha=0.2, linewidth=.5)
	    ax.plot(xe[i,:], ye[i,:], 'k-', alpha=0.2, linewidth=.5)
	ax.quiver(X,Y,dx,dy, angles='xy', edgecolor='k', facecolor='None', linewidth=.5, units='dots')
	ax.quiver(X,Y,dx,dy,R, alpha=.5, units='dots')
	ax.set_aspect('equal', adjustable='box')
	ax.set_xlabel('x (mm)')
	ax.set_ylabel('y (mm)')
	fig.colorbar(im, ax=ax, label='PSF centroid shift ($\\mu$as)')
	plt.savefig('./figures/UPDATED2_GD_full_field_zern_psd.pdf', bbox_inches='tight', transparent=True)
	plt.savefig('./figures/UPDATED2_GD_full_field_zern_psd.png', bbox_inches='tight', transparent=True)
	plt.tight_layout()
	plt.show()


if __name__ == '__main__':
	# plt.rc('text', usetex=True)
	# plt.rc('font', family='serif')

	# Loading in the WFEs
	wfe_p1 = load_variable('./data/compressed items/ADC_tWFE_p1_zern_20nm_psd_15nm.pkl')
	wfe_p2 = load_variable('./data/compressed items/ADC_tWFE_p2_zern_20nm_psd_15nm.pkl')
	wfe_p3 = load_variable('./data/compressed items/ADC_tWFE_p3_zern_20nm_psd_15nm.pkl')
	wfe_p4 = load_variable('./data/compressed items/ADC_tWFE_p4_zern_20nm_psd_15nm.pkl')
	
	

	# Rotation of ADC angles - [0.125, 0.25, 0.5, 0.75, 1, 2, 5, 10], 30, 45, 60, 90, 180]
	thetaList = [0.001, 0.125, 0.25, 0.5, 0.75, 1, 1.25, 1.5, 1.75, 2, 2.5, 3, 3.5, 4, 4.5, 5, 6, 7, 8, 9, 10]
	# Propagate through the field positions without any WFEs	
	for theta in thetaList:

		# Setup of simulation
		l = 1.49e-3
		Elist, Eplist = [], []
		xlist, ylist = [], []
		xplist, yplist = [], []
		x_expected_list, y_expected_list = [], []
		x_centroid_list, y_centroid_list = [], []

		# Set field points
		hxList, hyList = np.mgrid[-1:1:11j,-1:1:11j]-(1/300)
		hxList = hxList.flatten()
		hyList = hyList.flatten()

		for i in range(len(hxList)):
			hx = hxList[i]
			hy = hyList[i]
			print('Rotation: {:.3f} -  hx: {}	- 	hy: {}'.format(theta,hx,hy) )

			# Initializing MICADO paraxial model
			mcd = ParaxialMICADO(n=3001, m=3001, dpix=50, entrance_pupil='obscured')

			# Propagate through the MICADO system
			mcd.ELT_ADC.set_orientation(theta)
			E, x, y = mcd.propagate(l=l,surface=-1, zoom=False, show=False, 
									print_progress=False, hx=hx, hy=hy)

			# Find approximation for the centroid and amplitude as a first 
			# guess for the Obscured Airy pattern fit.
			amplitude_guess = np.max(np.abs(E))**2
			xbc, ybc, _, _ = mcd.centroid_shift(method='Barycenter')
			print('Barycenter guess: ', xbc, ybc)

			# Do the Obscured Airy Pattern Fit for the unperturbed PSF
			xe, ye, _, _ = mcd.centroid_shift(method='Obscured Airy', initialGuess=(amplitude_guess, xbc, ybc, 0.0, 0.03))
			print('Airy result: ', xe,  ye)
			x_expected_list.append(xe)
			y_expected_list.append(ye)

		

			# Load in the WFEs and set them to the corresponding ADC surfaces
			print('\nSetting the wavefront errors...')
			# null_wfe = np.zeros((3001,3001))
			mcd.ELT_ADC.set_orientation(0)
			mcd.ELT_ADC.p1.add_wavefront_error(wfe_p1)
			mcd.ELT_ADC.p2.add_wavefront_error(wfe_p2)
			mcd.ELT_ADC.p3.add_wavefront_error(wfe_p3)
			mcd.ELT_ADC.p4.add_wavefront_error(wfe_p4)


			# Propagate through the MICADO system
			print('Propagating with Wavefront Errors applied...')
			mcd.ELT_ADC.set_orientation(theta)
			Ep, xp, yp = mcd.propagate(l=l,surface=-1, zoom=False, show=False, 
									print_progress=False, hx=hx, hy=hy)

			# Do the Obscured Airy Pattern Fit for the perturbed PSF
			xf, yf, _, _ = mcd.centroid_shift(method='Obscured Airy', initialGuess=(amplitude_guess, xe, ye, 0.0, 0.03))

			# Resulting centroid shift due to the WFE
			x_centroid_list.append(xf)
			y_centroid_list.append(yf)
			R = np.sqrt( (xf-xe)**2 + (yf-ye)**2 ) * mcd.platescale * 1e6
			print('xf: {}, yf: {}, shift (uas): {}'.format(xf,yf, R))

			# Clear the memory by emptying the large arrays
			mcd = None
			wfe = None
			E, x, y = None, None, None
			Ep, xp, yp = None, None, None


		# Reshaping the arrays
		shape = (int(np.sqrt(len(x_expected_list))), int(np.sqrt(len(x_expected_list))))
		x_expected = np.asarray(x_expected_list).reshape(shape)
		y_expected = np.asarray(y_expected_list).reshape(shape)#.T

		# Saving the arrays
		save_variable('./data/ADC rotation/GD_obscured_PSF_xcentroid_expected_zern_psd_rot{}.pkl'.format(theta), x_expected)
		save_variable('./data/ADC rotation/GD_obscured_PSF_ycentroid_expected_zern_psd_rot{}.pkl'.format(theta), y_expected)


		# Reshaping the arrays
		shape = (int(np.sqrt(len(x_centroid_list))), int(np.sqrt(len(x_centroid_list))))
		x_centroid = np.asarray(x_centroid_list).reshape(shape)
		y_centroid = np.asarray(y_centroid_list).reshape(shape)
		
		# Saving the arrays
		save_variable('./data/ADC rotation/GD_obscured_PSF_xcentroid_found_with_wfe_airy_zern_psd_rot{}.pkl'.format(theta), x_centroid)
		save_variable('./data/ADC rotation/GD_obscured_PSF_ycentroid_found_with_wfe_airy_zern_psd_rot{}.pkl'.format(theta), y_centroid)

	# plot_distortions()