Added scripts folder with old scripts

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()
+
+
+
+
+
+
+
+
+
+

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()
+
+
+
+
+
+
+
+
+
+

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()
+
+	###########################################################

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()
+
+	###########################################################

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])
+	###########################################################
+
+
+
+
+
+
+
+
+

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()
+
+
+
+
+

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()
+
+	###########################################################

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()
+
+	###########################################################

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)))
+	
+
+

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()

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)))
+
+

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()
+
+
+
+
+

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()
+
+

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')
+
+
+
+
+
+
+
+
+
+
+
+	

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()
+	'''
+		
+
+
+
+
+
+
+
+
+
+
+	

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()
+	'''
+		
+
+
+
+
+
+
+
+
+
+
+	

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()
+
+
+
+
+
+
+
+
+

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()
+
+
+
+
+
+
+
+
+

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()
+
+
+
+
+
+
+
+

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()
+
+
+
+
+
+
+
+

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)))
+	
+
+

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()
+