Added features, fixed bugs and other small things

demonstration scripts/generate_PSF_with_WFE.py

@@ -0,0 +1,20 @@
+import numpy as np
+import sys
+sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics')
+from fourierOptics import *
+
+if __name__ == '__main__':
+	# Initiate MICADO paraxial model
+	mcd = ParaxialMICADO()
+
+	# Initiate WFE object
+	wfe = WavefrontError(mcd.n, mcd.m)
+
+	# Generate PSD surfaces
+	psd = wfe.generate_height_map(200e-6, 2, 1)
+
+	# Add it to ADC-P1
+	mcd.ELT_ADC.p1.add_wavefront_error(psd)
+
+	# Generate the PSF and plot it
+	E, x, y = mcd.propagate(l=1.35e-3,surface=-1, zoom=True, show=True)

demonstration scripts/generate_bahtinov_dispersion.py

@@ -0,0 +1,69 @@
+import numpy as np
+import sys
+sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics')
+from fourierOptics import *
+
+
+if __name__ == '__main__':
+	# Initiate MICADO paraxial model
+	mcd = ParaxialMICADO()
+	wfe = WavefrontError(mcd.n, mcd.m)
+
+	# Create the Bahtinov mask and set it on the entrance pupil (ELT-M1)
+	bahtinov = Profiles().bahtinov_mask(40, 10, mcd.n, mcd.m, dpix=0.33, N=30, oversample=1)
+	mcd.ELT_CAM.add_mask(bahtinov)
+
+	# Set up
+	Elist = []
+	lList = np.linspace(1.49e-3, 1.78e-3, 15).tolist()
+
+	# Dispersion of 40 mas over the band
+	dispMax = 40e-3 / 3600 		# maximum dispersion in degrees
+	dispersion = np.linspace(0, dispMax, len(lList)).tolist()
+
+	# Iterate over the wavelengths
+	for i in range(len(lList)):
+		l = lList[i]
+
+		# Generating the tilt to the input field.
+		d = dispersion[i]
+		field_tilt = wfe.generate_tilt(mcd.x,mcd.y,tx=0, ty=d)
+		
+		# Adding the tilt to the input field
+		Ein = Profiles().ELT_pupil(mcd.n,mcd.m,mcd.dpix).astype('complex128')\
+				* np.exp(-1j * 2*np.pi/l * field_tilt)
+		mcd.apply_mask(Ein, replace=True)
+
+		# Propagate the input field throug the paraxial model
+		E, x, y = mcd.propagate(l=l,surface=-1, zoom=False, show=False, 
+								image_output=[-4,4,-4,4], print_progress=True)
+		Elist.append(E)
+
+	# Average over all the wavelengths to get the final field.
+	Eca = np.asarray([np.abs(x) for x in Elist]).mean(axis=0) 	# Amplitude
+	Eci = Eca**2 												# Intensity
+	Ecp = np.asarray([np.angle(x) for x in Elist]).mean(axis=0) # Phase
+
+
+	# Convert the scale to sky-angles using the plate scale
+	x *= mcd.platescale*1000 			# x in milli arcseconds
+	y *= mcd.platescale*1000 			# y in milli arcseconds
+
+	# Plotting the results
+	fig = plt.figure(figsize=(8,3))
+	ax = fig.add_subplot(121)
+	im = ax.imshow(Eca, cmap='hot', origin='lower', extent=[x.min(),x.max(),y.min(),y.max()])
+	ax.set_xlim([x.min(),x.max()])
+	ax.set_ylim([y.min(),y.max()])
+	ax.set_xlabel('x [mas]')
+	ax.set_ylabel('y [mas]')
+	fig.colorbar(im, ax=ax, label='Amplitude')
+	ax2 = fig.add_subplot(122)
+	im2 = ax2.imshow(Ecp, cmap='hot', origin='lower', extent=[x.min(),x.max(),y.min(),y.max()])
+	ax2.set_xlim([x.min(),x.max()])
+	ax2.set_ylim([y.min(),y.max()])
+	ax2.set_xlabel('x [mas]')
+	ax2.set_ylabel('y [mas]')
+	fig.colorbar(im2, ax=ax2, label='Phase')
+	plt.tight_layout()
+	plt.show()

demonstration scripts/generate_field.py

@@ -0,0 +1,80 @@
+import numpy as np
+import sys
+sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics')
+from fourierOptics import *
+from pdb import set_trace
+
+def field_angle(hx, hy):
+	fx = hx * 0.0075
+	fy = hy * 0.0075
+	return fx, fy
+
+
+if __name__ == '__main__':
+	# Initializing MICADO paraxial model
+	mcd = ParaxialMICADO()
+	wfe = WavefrontError(mcd.n, mcd.m)
+
+	# Uncomment to load in WFE (coma)
+	# Problem here is that the field of view is too small to see 
+	# a radial effect...
+	# xin, yin = load_variable('../adc_data_coordinates.pkl')
+	# adcObj = ADC_surfaces(xin,yin)
+	# zern = wfe.zernike_cartesian([0,0,0,0,0,0,0,100e-5], adcObj.xz, adcObj.yz, rescale=False)
+	# mcd.ELT_ADC.p1.add_wavefront_error(zern)
+	
+	# Setup
+	l = 1.49e-3
+	Elist = []
+
+	# Set field points
+	hxList, hyList = np.mgrid[-0.03:0.03:3j,-0.03:0.03:3j]
+	hxList = hxList.flatten()
+	hyList = hyList.flatten()
+	
+	# Iterating over the field positions.
+	for i in range(len(hxList)):
+		hx = hxList[i]
+		hy = hyList[i]
+
+		# Calculate the relative wavefront tilt for the field position		
+		fieldx, fieldy = field_angle(hx,hy)
+		field_tilt = wfe.generate_tilt(mcd.x,mcd.y,tx=fieldx, ty=fieldy)
+
+		# Apply the field tilt
+		Ein = Profiles().ELT_pupil(mcd.n,mcd.m,mcd.dpix).astype('complex128')\
+				* np.exp(-1j * 2*np.pi/l * field_tilt)
+		mcd.apply_mask(Ein, replace=True)
+
+		# Propagate through the MICADO system
+		E, x, y = mcd.propagate(l=l,surface=-1, zoom=False, show=False, 
+								print_progress=False)
+		Elist.append(E)
+
+	# Average over all the wavelengths to get the final field.
+	Eca = np.asarray([np.abs(x) for x in Elist]).mean(axis=0) 	# Amplitude
+	Eci = Eca**2 												# Intensity
+	Ecp = np.asarray([np.angle(x) for x in Elist]).mean(axis=0) # Phase
+
+	# Convert the scale to sky-angles using the plate scale
+	x *= mcd.platescale*1000 			# x in milli arcseconds
+	y *= mcd.platescale*1000 			# y in milli arcseconds
+
+	# Plotting the results
+	fig = plt.figure(figsize=(8,3))
+	ax = fig.add_subplot(121)
+	im = ax.imshow(Eca, cmap='hot', origin='lower', extent=[x.min(),x.max(),y.min(),y.max()])
+	ax.set_xlim([x.min(),x.max()])
+	ax.set_ylim([y.min(),y.max()])
+	ax.set_xlabel('x [mas]')
+	ax.set_ylabel('y [mas]')
+	fig.colorbar(im, ax=ax, label='Amplitude')
+	ax2 = fig.add_subplot(122)
+	im2 = ax2.imshow(Ecp, cmap='hot', origin='lower', extent=[x.min(),x.max(),y.min(),y.max()])
+	ax2.set_xlim([x.min(),x.max()])
+	ax2.set_ylim([y.min(),y.max()])
+	ax2.set_xlabel('x [mas]')
+	ax2.set_ylabel('y [mas]')
+	fig.colorbar(im2, ax=ax2, label='Phase')
+	plt.tight_layout()
+	plt.show()

demonstration scripts/generate_propagation_figures.py

@@ -0,0 +1,16 @@
+import numpy as np
+import sys
+sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics')
+from fourierOptics import *
+import matplotlib.pyplot as plt
+from pdb import set_trace
+
+if __name__ == '__main__':
+	plt.rc('font', family='serif')
+	plt.rc('axes', labelsize='large')
+
+	# Initiate MICADO paraxial model
+	mcd = ParaxialMICADO()
+
+	# Generate the PSF and plot it
+	E, x, y = mcd.propagate(l=1.35e-3,surface=-1, zoom=True, show=True, save='all')

demonstration scripts/thesis_propagation_figures.py

@@ -0,0 +1,51 @@
+import numpy as np
+import sys
+sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics')
+from fourierOptics import *
+import matplotlib.pyplot as plt
+from pdb import set_trace
+
+if __name__ == '__main__':
+	plt.rc('font', family='serif')
+	plt.rc('axes', labelsize='large')
+
+	# Initiate MICADO paraxial model
+	mcd = ParaxialMICADO(n=3001, m=3001, dpix=50)
+
+	# Generate the PSF and plot it
+	# E, x, y = mcd.propagate(l=1.35e-3,surface=-1, zoom=True, show=True, save='all')
+
+	confList = [ ( 1, [-0.4, 0.4, -0.4, 0.4], 	'02_fourier_ELT-IFP.pdf'),
+				 ( 2, [-50, 50, -50, 50], 		'03_fourier_ELT-COL.pdf'),
+				 ( 4, [-50, 50, -50, 50], 		'04_fourier_ELT-MCD-ADC-P2.pdf'),
+				 (-1, [-1.2, 1.2, -1.2, 1.2], 	'05_fourier_ELT-MCD-DET.pdf')]
+
+
+	for i in range(len(confList)):
+		s = confList[i][0]
+		coordinates = confList[i][1]
+		savename = confList[i][2]
+		E, x, y = mcd.propagate(l=1.35e-3, surface=s, image_output=coordinates,
+								 zoom=False, show=False, save=savename)
+
+	Eout, x, y = mcd.propagate(l=1.35e-3, surface=0, zoom=False, show=False)
+	x /= 1e3 # convert to m
+	y /= 1e3 # convert to m
+	_fig = plt.figure(figsize=(8,3))
+	_ax = _fig.add_subplot(121)
+	_im = _ax.imshow(np.abs(Eout), cmap='hot', extent=[x.min(), x.max(), y.min(), y.max()])
+	_ax.set_xlim([-30, 30])
+	_ax.set_ylim([-30, 30])
+	_ax.set_xlabel('x [m]')
+	_ax.set_ylabel('y [m]')
+	_fig.colorbar(_im, ax=_ax, label='Amplitude')
+	_ax2 = _fig.add_subplot(122)
+	_im2 = _ax2.imshow(np.angle(Eout), cmap='hot', extent=[x.min(), x.max(), y.min(), y.max()])
+	_ax2.set_xlim([-30, 30])
+	_ax2.set_ylim([-30, 30])
+	_ax2.set_xlabel('x [m]')
+	_ax2.set_ylabel('y [m]')
+	_fig.colorbar(_im2, ax=_ax2, label='Phase')
+	plt.tight_layout()
+	plt.savefig('01_fourier_ELT-M1.pdf', bbox_inches='tight', transparent=True)
+	plt.close()#show()

fourierOptics/__init__.py

@@ -7,3 +7,4 @@ from .profiles import *
 from .psfFitter import *
 from .wavefrontErrors import *
 from .misc import *
+from .micado import *

fourierOptics/fourier.py

@@ -1,5 +1,4 @@
 import numpy as np
-from pdb import set_trace
 import matplotlib.pyplot as plt
 from cmath import sqrt as csqrt
 from scipy import interpolate
@@ -14,13 +13,14 @@ class FourierOptics:
 		self.l = l # Wavelength
 
 
-	def plane_wave_propagation(self, A, xin, yin, dz=None, oversample=None, show=False, save=False, scaleInputArray=False, outputCoordinates=None):
+	def plane_wave_propagation(self, A, xin, yin, dz=None, oversample=None, show=False, 
+								save=False, scaleInputArray=False, outputCoordinates=None):
 		'''
 			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:
+			Parameters:
 			A 					-	[array]	(Complex) Input field
 			dpix				- 	[float] pixel scale at which the calculations are done
 			dz 					- 	[float]	propagation distance
@@ -32,7 +32,7 @@ class FourierOptics:
 			outputCoordinates 	-	[list]	list the corners of the desired output coordinate plane.
 											Order: [xmin, xmax, ymin, ymax]
 
-			Output:
+			Returns:
 			E_out 				- 	[array] (Complex) Output field distribution
 			x, y 				-	[array] Output coordinates 
 		'''
@@ -59,9 +59,9 @@ class FourierOptics:
 		_n, _m = np.shape(A)
 		
 		# Calculate the pixel size from the xi, eta plane
-		_dpix = np.mean(np.diff(xin))
+		_dpix = np.mean(np.diff(xin.real))
 		if _dpix < 1e-99:
-			_dpix = np.mean(np.diff(xin))
+			_dpix = np.mean(np.diff(xin.real))
 
 		# A at z=0 is then
 		_E_xy0 		= pyfftw.empty_aligned((_n,_m), dtype='complex128')
@@ -93,8 +93,8 @@ class FourierOptics:
 
 		# Coordinates of E_xy0
 		# _y, _x = np.mgrid[-_n / 2 * _dpix : _n / 2 * _dpix : 1j*_n, -_m / 2 * _dpix : _m / 2 * _dpix : 1j * _m]
-		_y, _x = yin, xin
-
+		_y, _x = yin.real, xin.real
+		
 		if outputCoordinates is not None:
 			# Decide only use for plotting or can be used for output as well?
 			# Performing interpolation to get to the required final plane, done for real and imaginary part.
@@ -108,7 +108,6 @@ class FourierOptics:
 			_E_out_ip_imag	= _f_ip_imag(_xi, _yi)
 			_E_xyz 			= _E_out_ip_real + _E_out_ip_imag*1j
 			_y, _x 			= np.mgrid[outputCoordinates[0]:outputCoordinates[1]:1j*_n, outputCoordinates[2]:outputCoordinates[3]:1j * _m]
-			
 
 		_E_xyz_extent = [np.min(_x.real), np.max(_x.real), np.min(_y.real), np.max(_x.real)]
 
@@ -125,13 +124,13 @@ class FourierOptics:
 			im2 = ax2.imshow(np.angle(_E_xyz), cmap='hot', extent=_E_xyz_extent)
 			ax2.set_xlabel('x [mm]')
 			ax2.set_ylabel('y [mm]')
-			fig.colorbar(im2, ax=ax2)
+			fig.colorbar(im2, ax=ax2, label='Phase')
 			plt.tight_layout()
 		if save:
-			plt.savefig(save, bbox_inches='tight')
+			plt.savefig(save, bbox_inches='tight', transparent=True)
 		if show:
 			plt.show()
-		elif not show:
+		else:
 			plt.close()
 
 		return np.array([_E_xyz, _x.real, _y.real])
@@ -144,7 +143,7 @@ class FourierOptics:
 			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:
+			Parameters:
 			A 					-	[array]	(Complex) Input field
 			dpix				- 	[float] pixel scale at which the calculations are done
 			dz 					- 	[float]	propagation distance
@@ -158,7 +157,7 @@ class FourierOptics:
 											If 'auto': the output coordinates will be determined 
 											automatically.
 
-			Output:
+			Returns:
 			E_out 				- 	[array] (Complex) Output field distribution
 			x, y 				-	[array] Output coordinates 
 		'''
@@ -189,7 +188,7 @@ class FourierOptics:
 		_k 				= 2 * np.pi / self.l
 		
 		# Define (xi, eta) plane at z=0 in meters, with nxm points. We must use the pixel scale.
-		_xi, _eta = xin, yin
+		_xi, _eta = xin.real, yin.real
 		
 		
 		# Calculate the pixel size from the xi, eta plane
@@ -261,7 +260,7 @@ class FourierOptics:
 			_xip, _yip = _x, _y
 
 
-		if show:
+		if show or save:
 			# Plotting the results
 			_fig = plt.figure(figsize=(8,3))
 			_ax = _fig.add_subplot(121)
@@ -281,7 +280,10 @@ class FourierOptics:
 			plt.tight_layout()
 			if save:
 				plt.savefig(save, bbox_inches='tight', transparent=True)
-			plt.show()
+			if show:
+				plt.show()
+			else:
+				plt.close()
 
 		_x, _y = np.meshgrid(_xip, _yip)
 		return np.array([_E_out_ip, _x.real, _y.real])

fourierOptics/micado.py

@@ -0,0 +1,438 @@
+import numpy as np
+from .fourier import *
+from .gaussian import *
+from .opticalTrain import *
+from .optics import *
+from .profiles import *
+from .psfFitter import *
+from .wavefrontErrors import *
+from .misc import *
+
+
+
+
+class ParaxialMICADO:
+	def __init__(self, l=1.35e-3, n=1001, m=1001, dpix=160):
+		self.l 				= l
+		self.dpix 			= dpix
+		self.n, self.m 		= n, m
+		self.x, self.y 		= np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m))	
+		self.ELT_MASK 		= Profiles().ELT_pupil(n,m,dpix).astype('complex128')
+
+		# Built optical system elements
+		self.ELT_MIRROR 	= ParabolicMirror(dx=0, dy=0, name='ELT-paraxial', CA=38542, f=6.84022e5)	
+		self.ELT_FP			= OpticalSurface(name='MCD-IFP')
+		self.ELT_COL		= ParaxialLens(CA=154.54, f=1418.94, name='MCD-COL')
+		self.ELT_ADC 		= ADCObject()
+		self.ELT_CAM		= ParaxialLens(CA=89.2, f=4278.32701, name='MCD-CAM')
+		self.ELT_DET		= OpticalSurface(name='MCD-DET')
+
+		# Built optical train
+		self.MCD 			= OpticalTrain()
+		self.MCD.add_object(self.ELT_MIRROR, 	dist=0)
+		self.MCD.add_object(self.ELT_FP, 		dist=self.ELT_MIRROR.f)
+		self.MCD.add_object(self.ELT_COL, 		dist=self.ELT_COL.f)
+		self.ELT_ADC.add_to_optical_train(self.MCD, 1209.4718)
+		self.MCD.add_object(self.ELT_CAM, 		dist=50)
+		self.MCD.add_object(self.ELT_DET, 		dist=self.ELT_CAM.f)
+
+		self.last_PSF 		= None
+		self.clean_PSF 		= None
+		self.platescale 	= self.set_plate_scale()
+		print ('Your resolution is {}x{} with a pixel size of {}'.format(self.n,self.m,self.dpix))
+		print('This means your grid is {} mm by {} mm, with each pixel {} mm.'.format(
+			self.x.max()-self.x.min(), self.y.max()-self.y.min(), self.dpix))
+
+
+	def apply_mask(self, mask, replace=True):
+		'''
+			Set a mask to the ELT entrance pupil.
+		'''
+		if replace:
+			self.ELT_MASK = mask.astype('complex128')
+		elif not replace:
+			self.ELT_MASK *= mask
+
+
+	def set_PSF(self, PSF, x, y, perfect=True):
+		'''
+			If calculating another PSF takes to long,
+			one can set a custom PSF using this 
+			function.
+
+			Can be set as the 'clean' PSF or as another
+			'last' PSF, with the perfect flag.
+		'''
+		if perfect:
+			self.clean_PSF = (PSF, x, y)
+		elif not perfect:
+			self.last_PSF  = (PSF, x, y)
+
+
+	def centroid_shift(self, method='Coupling'):
+		'''
+			Do a centroid shift calculation for the given method.
+			This is a quick way of doing it. It only returns the 
+			centroid shift (and uncertainties) and no other parameters.
+		'''
+		assert (self.last_PSF is not None), 'PSF necessary for fitting.'
+		
+		PSF, x, y = self.last_PSF
+		fitter = PSFFitter(x.real,y.real,np.abs(PSF)**2)
+
+		if method == 'Coupling':
+			assert (self.clean_PSF is not None), 'Perfect PSF necessary for comparison.'
+			Eclean, xclean, yclean = self.clean_PSF
+			fitted_PSF, popt, var, coupling = fitter.fit_custom_field(Eclean, PSF, return_coupling=True)
+			x0, y0 		= np.asarray(popt).T
+			sx0, sy0 	= np.asarray(np.sqrt(var)).T
+			return x0, y0, sx0, sy0, coupling
+
+		elif method == 'Airy':
+			fitted_PSF, popt, pcov = fitter.fit_airy_PSF(initialGuess=(0.03**2, 0.0, 0.0, 0.0, 0.03))
+			_, x0, y0, _, _ = popt_airy
+			pcov_list = [ np.sqrt(pcov[i][i]) for i in range(len(pcov)) ]
+			_, sx0, sy0, _, _ = np.asarray(pcov_list).T
+
+		elif method == 'Gaussian':
+			fitted_PSF, popt, pcov = fitter.fit_gaussian_PSF(initialGuess=(0.04, 0, 0, 0.01, 0.01, 0, 0))
+			_, x0, y0, _, _, _, _ = popt
+			pcov_list = [ np.sqrt(pcov[i][i]) for i in range(len(pcov)) ]
+			_, sx0, sy0, _, _, _, _ = np.asarray(pcov_list).T
+
+		elif method == 'Moffat':
+			fitted_PSF, popt, pcov = fitter.fit_moffat_PSF(initialGuess=(0.04, 0, 0, 0.01, 0.01, 0, 0, 4))
+			_, x0, y0, _, _, _, _, _ = popt
+			pcov_list = [ np.sqrt(pcov[i][i]) for i in range(len(pcov)) ]
+			_, sx0, sy0, _, _, _, _, _ = np.asarray(pcov_list).T
+
+		return x0, y0, sx0, sy0
+
+
+	def propagate(self, l=None, surface=-1, zoom=False, load_up_to_adc=None, coordinates=None, 
+		show=False, image_output=None, save=None, apply_interpolation=True, print_progress=True):
+		'''
+			Propagate the input wavefront through the MICADO paraxial
+			model. Optionally the propagation starts at the ADC, if 
+			a data file with the wavefront at that point is defined.
+
+			Parameters:
+			l 				- Desired wavelength of the propagation
+			surface 		- Surface to propagate to. Default: -1 = Detector plane
+			zoom 			- Zoom in on the PSF on the detector?
+			load_up_to_adc 	- Use a .pkl file to load in the wavefront just before the adc.
+			coordinates 	- User defined coordinates that define the propagation sampling.
+			show_final 		- Plots the final amplitude and phase if True.
+
+			Returns:
+			Exyz 			- Complex field with shape n,m describing the amplitude and phase
+							  of the wavefront at the output surface.
+			x, y 			- Corresponding x and y coordinates, in mm.
+		'''
+		# Set up wavelength
+		l = self.l if not l else l
+	
+		# Setting up the arguments for each surface propgation
+		surfaces 			= self.get_surfaces()
+		surface 			= len(surfaces)-1 if surface == -1 else surface
+		coordinatesList 	= self.get_coordinates(coordinates)
+		show,show_final_zoom= self.get_show_image(show, zoom, len(surfaces), surface)
+		output_coordinates 	= self.get_image_coordinates(image_output, coordinatesList, len(surfaces), surface)
+		saveList 			= self.get_save_image(save, len(surfaces), surface)
+		apply_ip 			= self.get_image_interpolation(apply_interpolation, zoom, len(surfaces), surface)
+
+		# Set up the FourierOptics object for propagation
+		fou = FourierOptics(l=l)
+		x, y = self.x, self.y
+		p = 1 # Used as counter for the ADC prisms.
+		
+		# Load the data up to the ADC (useful for many iterations on the same 
+		# calculation, where differences only appear at or after the ADC)
+		if load_up_to_adc is not None:
+			Eout, x, y = load_variable(load_up_to_adc)
+			x = x.real
+			y = y.real
+			s = 4
+		elif load_up_to_adc is None:
+			s = 0
+
+		# Start propgation
+		if print_progress: print('Starting propagation...')
+		while s <= surface:
+			if (surfaces[s] == self.ELT_MASK).all():
+				Ein  = surfaces[s]
+				Eout = surfaces[s]
+				if show[s] or saveList[s]:
+					_fig = plt.figure(figsize=(8,3))
+					_ax = _fig.add_subplot(121)
+					_im = _ax.imshow(np.abs(Eout), cmap='hot', extent=[x.min(), x.max(), y.min(), y.max()])
+					_ax.set_xlim([x.min(), x.max()])
+					_ax.set_ylim([y.min(), y.max()])
+					_ax.set_xlabel('x [mm]')
+					_ax.set_ylabel('y [mm]')
+					_fig.colorbar(_im, ax=_ax, label='Amplitude')
+					_ax2 = _fig.add_subplot(122)
+					_im2 = _ax2.imshow(np.angle(Eout), cmap='hot', extent=[x.min(), x.max(), y.min(), y.max()])
+					_ax2.set_xlim([x.min(), x.max()])
+					_ax2.set_ylim([y.min(), y.max()])
+					_ax2.set_xlabel('x [mm]')
+					_ax2.set_ylabel('y [mm]')
+					_fig.colorbar(_im2, ax=_ax2, label='Phase')
+					plt.tight_layout()
+					if saveList[s]:	
+						plt.savefig(saveList[s], bbox_inches='tight', transparent=True)
+					if show[s]:
+						plt.show()
+					else:
+						plt.close()
+
+			
+			elif surfaces[s] == self.ELT_MIRROR:
+				if print_progress: print('Propagating from ELT-M1 to ELT-IFP')
+				Ein = Eout * self.ELT_MIRROR.phase_delay(x,y,l,paraxial=True)
+				if s == surface and zoom:
+					xzoom, yzoom = x, y
+					Einzoom = Ein
+					fzoom = self.ELT_MIRROR.f
+				Eout, x, y = fou.fresnel_propagation(Ein, x, y, dz=self.ELT_MIRROR.f, oversample=None, 
+					show=show[s], save=saveList[s], outputCoordinates=output_coordinates[s], 
+					apply_interpolation=apply_ip[s])
+
+			elif surfaces[s] == self.ELT_FP:
+				if print_progress: print('Propagating from ELT-IFP to MCD-COL')
+				Ein = Eout
+				Eout, x, y = fou.fresnel_propagation(Ein, x, y, dz=self.ELT_COL.f, oversample=None, 
+					show=show[s], save=saveList[s], outputCoordinates=output_coordinates[s], 
+					apply_interpolation=apply_ip[s])
+
+			elif surfaces[s] == self.ELT_COL:
+				if print_progress: print('Propagating from MCD-COL to MCD-ADC-P1')
+				Ein = Eout * self.ELT_COL.phase_delay(x,y,l, paraxial=True)
+				Eout, x, y = fou.plane_wave_propagation(Ein, x, y, dz=self.MCD.sequentialDistance[3], oversample=None,
+					show=show[s], save=saveList[s], outputCoordinates=output_coordinates[s])
+
+			elif surfaces[s] == self.ELT_ADC and p == 1:
+				if print_progress: print('Propagating from MCD-ADC-P1 to MCD-ADC-P2')
+				Ein = Eout * self.ELT_ADC.get_phase_delay(x,y,l, theta=None, paraxial=True, prism=1)
+				Eout, x, y = fou.plane_wave_propagation(Ein, x, y, dz=self.MCD.sequentialDistance[4], oversample=None,
+					show=show[s], save=saveList[s], outputCoordinates=output_coordinates[s])	
+				p += 1
+
+			elif surfaces[s] == self.ELT_ADC and p == 2:
+				if print_progress: print('Propagating from MCD-ADC-P2 to MCD-ADC-P3')
+				Ein = Eout * self.ELT_ADC.get_phase_delay(x,y,l, theta=None, paraxial=True, prism=2)
+				Eout, x, y = fou.plane_wave_propagation(Ein, x, y, dz=self.MCD.sequentialDistance[5], oversample=None,
+					show=show[s], save=saveList[s], outputCoordinates=output_coordinates[s])	
+				p += 1
+
+			elif surfaces[s] == self.ELT_ADC and p == 3:
+				if print_progress: print('Propagating from MCD-ADC-P3 to MCD-ADC-P4')
+				Ein = Eout * self.ELT_ADC.get_phase_delay(x,y,l, theta=None, paraxial=True, prism=3)
+				Eout, x, y = fou.plane_wave_propagation(Ein, x, y, dz=self.MCD.sequentialDistance[6], oversample=None,
+					show=show[s], save=saveList[s], outputCoordinates=output_coordinates[s])	
+				p += 1
+
+			elif surfaces[s] == self.ELT_ADC and p == 4:
+				if print_progress: print('Propagating from MCD-ADC-P4 to MCD-CAM')
+				Ein = Eout * self.ELT_ADC.get_phase_delay(x,y,l, theta=None, paraxial=True, prism=4)
+				Eout, x, y = fou.plane_wave_propagation(Ein, x, y, dz=self.MCD.sequentialDistance[7], oversample=None,
+					show=show[s], save=saveList[s], outputCoordinates=output_coordinates[s])	
+				p += 1
+
+			elif surfaces[s] == self.ELT_DET:
+				if print_progress: print('Propagating from MCD-CAM to MCD-DET')
+				Ein = Eout * self.ELT_CAM.phase_delay(x,y,l, paraxial=True)
+				if s == surface and zoom:
+					xzoom, yzoom = x.real, y.real
+					Einzoom = Ein
+					fzoom = self.ELT_CAM.f
+				Eout, x, y = fou.fresnel_propagation(Ein, x, y, dz=self.ELT_CAM.f, oversample=None, 
+				show=show[s], save=saveList[s], outputCoordinates=output_coordinates[s], apply_interpolation=apply_ip[s])
+			s += 1
+
+		if zoom:
+			if print_progress: print('Zooming in on the PSF location...')
+			imax_1, imax_2 = find_maximum(np.abs(Eout))
+			xmax = x[imax_1, imax_2]
+			ymax = y[imax_1, imax_2]
+			ind = s-3 if (s-3)>0 else 0
+			newCoordinates = np.asarray(coordinatesList[ind]) + np.array([xmax, xmax, ymax, ymax])
+			newCoordinates = newCoordinates.real.tolist()
+			Eout, x, y = fou.fresnel_propagation(Einzoom, xzoom, yzoom, dz=fzoom, oversample=None, 
+				show=show_final_zoom, save=saveList[s-1], outputCoordinates=newCoordinates, apply_interpolation=True)
+		
+		if print_progress: print('Propagation finished!')
+		self.last_PSF = (Eout, x, y)
+		return Eout, x, y
+
+
+	def get_surfaces(self):
+		_surfaces = [	
+					self.ELT_MASK, 
+					self.ELT_MIRROR,
+					self.ELT_FP, 
+					self.ELT_COL, 
+					self.ELT_ADC,
+					self.ELT_ADC,
+					self.ELT_ADC,
+					self.ELT_ADC,
+					self.ELT_CAM,
+					self.ELT_DET
+					]
+		return _surfaces
+
+
+	def get_coordinates(self, coordinates):
+		# Output coordinates:
+		if coordinates is None:
+			_coordinatesList = [
+								[-1.0, 1.0, -1.0, 1.0],
+								[-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]
+							  ]
+		elif coordinates is not None:
+			_coordinatesList = coordinates
+		return _coordinatesList
+
+	def get_show_image(self, show, zoom, N, surf):
+		# Set show values. Will only show the image plane.
+		_showlist = [False] * N
+		_showlist[surf] = show
+		_show_final_zoom = False
+		if zoom:
+			_show_final_zoom = show
+			_showlist[surf] = False
+		_show = _showlist
+		return _show, _show_final_zoom
+
+	def get_image_coordinates(self, image_output, coordinates, N, surf):
+		# Set the image coordinates
+		_output_coordinates = [None] * N
+		_output_coordinates[2] = coordinates[1]
+		if image_output is not None:
+			_output_coordinates[surf] = image_output
+		return _output_coordinates
+
+	def get_save_image(self, save, N, surf):
+		# Set the save values
+		_savelist = [False] * N
+		if save == 'all':
+			_savelist = [
+						'01_ELT_MASK_output.png',
+						'02_ELT_FP_output.png',
+						'03_ELT_MCD_COL_output.png',
+						'04_ELT_MCD_ADC_P1_output.png',
+						'05_ELT_MCD_ADC_P2_output.png',
+						'06_ELT_MCD_ADC_P3_output.png',
+						'07_ELT_MCD_ADC_P4_output.png',
+						'08_ELT_MCD_CAM_output.png',
+						'09_ELT_MCD_DET_output.png',
+						'10_ELT_MCD_DET_output_zoomed.png'
+						]
+		elif save is not None:
+			_savelist[surf] = save
+		return _savelist
+
+	def get_image_interpolation(self, apply_interpolation, zoom, N, surf):
+		# Apply interpolation to the output?
+		_apply_ip = [False] * N
+		if not zoom and apply_interpolation and surf!=2:
+			_apply_ip[surf] = True
+		return _apply_ip
+
+	def set_plate_scale(self):
+		_ps = (3600*180/np.pi) / (self.ELT_MIRROR.CA * 53.4466)
+		self.platescale = _ps
+		return _ps
+
+
+
+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

fourierOptics/misc.py

@@ -90,6 +90,4 @@ def load_variable(filename):
 	'''
 	with open(filename, 'rb') as f:
 		variable = pickle.load(f)
-	return variable
-
-
+	return variable

fourierOptics/optics.py

@@ -14,12 +14,15 @@ class LensObject:
 		self.f 			= None
 		self.sag 		= None
 		self.wavefrontError = False
+		self.mask 		= None
 
 	def add_wavefront_error(self, z):
 		assert len(z.shape)==2, 'Wavefront data is not 2D. Cannot continue'
 		self.sag = z
 		self.wavefrontError = True
 
+	def add_mask(self, mask):
+		self.mask = mask
 
 
 
@@ -42,6 +45,8 @@ class ParaxialLens(LensObject):
 		if self.wavefrontError:
 			_t_l *=  np.exp(-1j * _k * self.sag / l)
 
+		if self.mask is not None:
+			_t_l *= self.mask
 		# Masking out the parts that fall outside the lens clear aperture
 		# dist_from_center = np.sqrt(x**2 + y**2)
 		# _lensMask = (dist_from_center <= (self.CA/2)).astype(int)
@@ -152,6 +157,9 @@ class SphericalLens(LensObject):
 		if self.wavefrontError:
 			_t_l *=  np.exp(-1j * _k * (_n - 1) * self.sag)
 
+		if self.mask is not None:
+			_t_l *= self.mask
+
 		# Masking out the parts that fall outside the lens clear aperture
 		dist_from_center = np.sqrt(x**2 + y**2)
 		_lensMask = (dist_from_center <= (self.CA/2)).astype(int)
@@ -291,6 +299,9 @@ class EvenAsphereLens(LensObject):
 		if self.wavefrontError:
 			_t_l *=  np.exp(-1j * _k * (_n - 1) * self.sag)
 
+		if self.mask is not None:
+			_t_l *= self.mask
+
 		# Masking out the parts that fall outside the lens clear aperture
 		dist_from_center = np.sqrt(x**2 + y**2)
 		_lensMask = (dist_from_center <= (self.CA/2)).astype(int)
@@ -317,13 +328,17 @@ class MirrorObject:
 		self.type 			= None
 		self.sag 			= None
 		self.wavefrontError = False
-
+		self.mask 			= None
 
 	def add_wavefront_error(self, z):
 		assert len(z.shape)==2, 'Wavefront data does not have the proper shape. Cannot continue'
 		self.sag = z
 		self.wavefrontError = True
 
+	def add_mask(self, mask):
+		self.mask = mask
+
+
 
 
 
@@ -346,6 +361,9 @@ class ParabolicMirror(MirrorObject):
 		else:
 			_t_m = np.exp(1j * _k * 2 * _dz)
 
+		if self.mask is not None:
+			_t_m *= self.mask
+
 		# Masking out the parts that fall outside the mirror clear aperture
 		dist_from_center = np.sqrt(x**2 + y**2)
 		_mirrorMask = (dist_from_center <= (self.CA/2)).astype(int)
@@ -375,6 +393,9 @@ class SphericalMirror(MirrorObject):
 		else:
 			_t_m = np.exp(1j * _k * 2 * _dz)
 
+		if self.mask is not None:
+			_t_m *= self.mask
+
 		# Masking out the parts that fall outside the mirror clear aperture
 		dist_from_center = np.sqrt(x**2 + y**2)
 		_mirrorMask = (dist_from_center <= (self.CA/2)).astype(int)
@@ -400,6 +421,10 @@ class FreeformMirror(MirrorObject):
 			_t_m = np.exp(1j * _k * 2 * ( self.z + self.sag) )
 		else:
 			_t_m = np.exp(1j * _k * 2 * self.z)
+
+		if self.mask is not None:
+			_t_m *= self.mask
+
 		return _t_m
 
 
@@ -415,9 +440,17 @@ class OpticalSurface:
 		if name:
 			assert type(name)==str, 'No string for a name.'
 			self.name = name
+		self.mask = None
+
+	def add_mask(self, mask):
+		self.mask = mask
 
 	def phase_delay(self, x, y, l, **kwargs):
-		return np.ones(np.shape(x))
+		_t_l = np.ones(np.shape(x))
+
+		if self.mask is not None:
+			_t_l *= self.mask
+		return _t_l
 
 
 
@@ -639,6 +672,7 @@ class Prism:
 		self.sag2		= None
 		self.wavefrontError = False
 		self.thetaPrism	= 0
+		self.mask 		= None
 
 		# Remember original tilts for reset:
 		self.tilt1x_OG	= tilt1
@@ -709,6 +743,9 @@ class Prism:
 			self.sag2 = z
 		self.wavefrontError = True
 
+	def add_mask(self, mask):
+		self.mask = mask
+
 
 	def phase_delay(self, x, y, l, theta=None, paraxial=True):
 		'''
@@ -748,6 +785,8 @@ class Prism:
 				elif self.sag2 is not None:
 					_t_p *=  np.exp(-1j * _k * (_n - 1) * self.sag2)
 
+			if self.mask is not None:
+				_t_p *= self.mask
 			# _t_p *= _prismMask
 			return _t_p
 

fourierOptics/profiles.py

@@ -70,4 +70,51 @@ class Profiles:
 		_x				= np.linspace(-n/2*dpix, n/2*dpix, n)
 		_y				= np.linspace(-m/2*dpix, m/2*dpix, m)
 		_ELT_pupil 		= _f_pupil(_x, _y)
-		return _ELT_pupil
+		return _ELT_pupil
+
+
+	def bahtinov_mask(self, r1, r2, n, m, dpix, N, oversample=1, half_width_spiders=None):
+		'''
+			Creates a Bahtinov mask profile. 
+			This profile is a popular mask to focus a telescope in the 
+			amateur community.
+		'''
+		assert r1 > r2, "r2 < r1, change input"
+
+		if half_width_spiders is None:
+			half_width_spiders = 0.03 * r1
+		
+		# Obscured top hat
+		j,i = np.mgrid[-n/2*dpix : n/2*dpix : 1j*n*oversample, -m/2*dpix : m/2*dpix : 1j*m*oversample]
+		dist_from_center = np.sqrt(i**2 + j**2)
+		_Anew = (dist_from_center <= (r1)) & (dist_from_center >= (r2))
+		_Anew =  _Anew.astype(int)
+
+		# Three quadrants of Bahtinov mask
+		_M1 = (i < -half_width_spiders)
+		_M2 = (i > half_width_spiders) & (j > half_width_spiders)
+		_M3 = (i > half_width_spiders) & (j < -half_width_spiders)
+		_MC = _M1 | _M2 | _M3    
+
+		# Generating the lines
+		_d  =  2 * r1 / int(N)
+		_M4 = step(j, i, period=_d, rotation=0)
+		_M5 = step(j, i, period=_d, rotation=30)
+		_M6 = step(j, i, period=_d, rotation=-30)
+		_MC = (_M1 & _M4) | (_M2 & _M5) | (_M3 & _M6) 
+
+		# Combining with the obscured top hat
+		_Anew = _Anew * _MC
+		return _Anew.astype(int) 
+
+
+def step(x, y, period=1, rotation=0):
+    '''
+		Function that returns a 2D square wave with period and rotation.
+    '''
+    theta = np.deg2rad(rotation)
+    xn = x * np.cos(theta) - y * np.sin(theta) 
+    yn = x * np.sin(theta) + y * np.cos(theta)
+    a = ((2/period * xn ).astype(int) % 2) & (xn>= 0)
+    b = ((2/period * xn - 1).astype(int) % 2) & (xn< 0)
+    return a + b