Initial commit

.gitignore

@@ -1,30 +1,6 @@
-# Byte-compiled / optimized / DLL files
+.DS_Store
-__pycache__/
+__pychache__/
-*.py[cod]
+.ipynb_checkpoints/
-*$py.class
-
-# C extensions
-*.so
-
-# Distribution / packaging
-.Python
-build/
-develop-eggs/
-dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-wheels/
-pip-wheel-metadata/
-share/python-wheels/
-*.egg-info/
-.installed.cfg
-*.egg
 MANIFEST
 
 # PyInstaller

fourierOptics/__init__.py

@@ -0,0 +1,9 @@
+import os
+from .fourier import *
+from .gaussian import *
+from .opticalTrain import *
+from .optics import *
+from .profiles import *
+from .psfFitter import *
+from .wavefrontErrors import *
+from .misc import *

fourierOptics/data/glassCatalog.py

@@ -0,0 +1,12 @@
+#!/usr/bin/env Python3
+
+# Pyton dictionary with some glass types
+
+catalog = {
+	
+	'N-BK7' 	: [1.039612120E+00, 6.000698670E-03, 2.317923440E-01, 2.001791440E-02, 1.010469450E+00, 1.035606530E+02	],
+	'SiO2'		: [0.6961663, 		0.0684043, 		 0.4079426, 	  0.1162414, 	   0.8974794, 		9.896161		],
+	'F2'		: [1.345333590E+00, 9.977438710E-03, 2.090731760E-01, 4.704507670E-02, 9.373571620E-01, 1.118867640E+02	],
+	'S-FPL51'	: [1.170105050E+00, 6.162039240E-03, 4.757107830E-02, 2.633728760E-02, 7.638324450E-01, 1.418826420E+02	],
+	'S-LAH71'	: [1.982800310E+00, 1.189874560E-02, 3.167584500E-01, 5.271560010E-02, 2.444726460E+00, 2.132206970E+02	],
+}

fourierOptics/fourier.py

@@ -0,0 +1,440 @@
+import numpy as np
+from pdb import set_trace
+import matplotlib.pyplot as plt
+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
+
+
+class FourierOptics:
+	def __init__(self, l):
+		self.l = l # Wavelength
+
+
+	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:
+			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)
+		
+		# Calculate the pixel size from the xi, eta plane
+		_dpix = np.mean(np.diff(xin))
+		if _dpix < 1e-99:
+			_dpix = np.mean(np.diff(xin))
+
+		# 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)
+		_kz 		 = self.phase_exponent(_fx, _fy, self.l) 	# origin at start of array
+		_phaseFactor = np.exp(-1j * _kz * dz) 					# origin at start of array
+		
+		# Calculating the spatial frequencies in k-space
+		_kx 		= np.fft.fftshift( _fx * 2 * np.pi )
+		_ky 		= np.fft.fftshift( _fy * 2 * np.pi )
+
+		_E_xyz 		= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_A_in 		= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_A_in[:] 	= _Az0 * _phaseFactor
+
+		_FFTW_Obj2 	= pyfftw.FFTW( _A_in, _E_xyz, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
+		_FFTW_Obj2()
+		_E_xyz 		= np.fft.fftshift(_E_xyz) 				# use fftshift to move the origin from edge to center
+		_E_xyz /= np.sqrt(_IntensityNormFactor) 			# Normalization back to initial energy.
+
+		# 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
+
+		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.
+			_x				= np.linspace(-_n/2*_dpix, _n/2*_dpix, _n)
+			_y				= np.linspace(-_m/2*_dpix, _m/2*_dpix, _m)
+			_f_ip_real		= interpolate.interp2d(_x, _y, _E_xyz.real, kind='linear')
+			_f_ip_imag 		= interpolate.interp2d(_x, _y, _E_xyz.imag, kind='linear')
+			_xi				= np.linspace(outputCoordinates[0],outputCoordinates[1],_n)
+			_yi 			= np.linspace(outputCoordinates[2],outputCoordinates[3],_m)
+			_E_out_ip_real 	= _f_ip_real(_xi, _yi)
+			_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)]
+
+
+
+		if show or save:
+			fig = plt.figure(figsize=(8,3))
+			ax = fig.add_subplot(121)
+			im = ax.imshow(np.abs(_E_xyz), cmap='hot', extent=_E_xyz_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_xyz), cmap='hot', extent=_E_xyz_extent)
+			ax2.set_xlabel('x [mm]')
+			ax2.set_ylabel('y [mm]')
+			fig.colorbar(im2, ax=ax2)
+			plt.tight_layout()
+		if save:
+			plt.savefig(save, bbox_inches='tight')
+		if show:
+			plt.show()
+		elif not show:
+			plt.close()
+
+		return np.array([_E_xyz, _x.real, _y.real])
+
+
+	def fresnel_propagation(self, A, xin, yin, dz=None, oversample=None, show=False, save=False, 
+				scaleInputArray=False, outputCoordinates=None, apply_interpolation=True):
+		'''
+			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]
+											If 'auto': the output coordinates will be determined 
+											automatically.
+
+			Output:
+			E_out 				- 	[array] (Complex) Output field distribution
+			x, y 				-	[array] Output coordinates 
+		'''
+		assert np.shape(xin) == np.shape(A), 'Input coordinates (xin) do not match shape of input field.'
+		assert np.shape(yin) == np.shape(A), 'Input coordinates (yin) do not match shape of input field.'
+		A = np.conj(A)
+
+		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='linear')
+			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)
+		_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
+		
+		
+		# Calculate the pixel size from the xi, eta plane
+		_dpix = np.mean(np.diff(_xi))
+		if _dpix < 1e-99:
+			_dpix = np.mean(np.diff(_xi.T))
+
+		# Gaussian test: use 8*w0 radius for FFT:
+		if outputCoordinates == 'auto':
+			_w0_fit, _ 		= leastsq(self.minimize, self.l, args=(0, _xi, _eta, A))
+			_gauss 			= GaussianBeam(_w0_fit, self.l)
+			_halfwidth 		= 8 * _gauss.beam_radius(dz)[0]
+			_xip			= np.linspace(-_halfwidth, _halfwidth, _n) 						# Desired x,y - plane
+			_yip			= np.linspace(-_halfwidth, _halfwidth, _m) 						# Desired x,y - plane
+		elif outputCoordinates is not None:
+			_xip			= np.linspace(outputCoordinates[0], outputCoordinates[1], _n) 	# Desired x,y - plane
+			_yip			= np.linspace(outputCoordinates[2], outputCoordinates[3], _m) 	# Desired x,y - plane
+		else:
+			_xip = xin[0]
+			_yip = yin[:,0]
+
+
+		# The quadratic phase exponential of the input wavefront
+		_phaseComponent = np.exp(1j * _k / (2 * dz) * (_xi**2 + _eta**2))
+
+		# Complete input
+		_U_in 			= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_U_in[:]		= A * _phaseComponent
+		
+		# Output of the FFT evaluated at frequency fx = x / (lambda * z)
+		_U_out 			= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_FFTW_Obj		= pyfftw.FFTW( np.fft.ifftshift(_U_in), _U_out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
+		_FFTW_Obj()
+		_U_out 			= np.fft.fftshift(_U_out)
+
+		# Evaluate frequencies:
+		_fx 			= np.fft.fftshift(np.fft.fftfreq(np.shape(_xi)[0], d=_dpix))
+		_fy 			= np.fft.fftshift(np.fft.fftfreq(np.shape(_xi)[1], d=_dpix))
+		_x 				= _fx * self.l * dz
+		_y 				= _fy * self.l * dz
+		_x2, _y2 		= np.meshgrid(_x, _y)
+
+
+		# Final output and frequencies
+		_pre_amble 		= np.exp(1j * _k * dz) / (1j * self.l * dz) * np.exp(1j * _k / (2 * dz) * (_x2**2 + _y2**2))
+		_E_out 			= _pre_amble * _U_out
+
+		# Normalization of the output values
+		_IntensityNormFactor = np.sum(np.abs(A)**2)
+		_E_out /= (np.abs(np.exp(1j*_k*dz)/(1j*self.l*dz)) * np.sqrt(_IntensityNormFactor) * np.sqrt(_n*_m))
+
+		# Performing interpolation to get to the required final plane, done for real and imaginary part.
+		if outputCoordinates and apply_interpolation:
+			_f_ip_real		= interpolate.interp2d(_x.real, _y.real, _E_out.real, kind='linear')
+			_f_ip_imag 		= interpolate.interp2d(_x.real, _y.real, _E_out.imag, kind='linear')
+			_E_out_ip_real 	= _f_ip_real(_xip, _yip)
+			_E_out_ip_imag	= _f_ip_imag(_xip, _yip)
+			_E_out_ip 		= _E_out_ip_real + _E_out_ip_imag*1j
+			_E_out_ip 		/= np.sqrt(((_x.max()/_xip.max())*(_y.max()/_yip.max())))
+			_E_out_ip = np.conj(_E_out_ip)
+		elif outputCoordinates and not apply_interpolation:
+			_arraySelection = self.select_subsection(_x, _y, outputCoordinates)
+			_xip 			= _x.real[_arraySelection[0]:_arraySelection[1]+1]
+			_yip 			= _y.real[_arraySelection[2]:_arraySelection[3]+1]
+			_E_out_ip 		= _E_out[_arraySelection[0]:_arraySelection[1]+1,_arraySelection[2]:_arraySelection[3]+1]
+			_E_out_ip = np.conj(_E_out_ip)
+		else:
+			_E_out_ip = _E_out
+			_xip, _yip = _x, _y
+
+
+		if show:
+			# Plotting the results
+			_fig = plt.figure(figsize=(8,3))
+			_ax = _fig.add_subplot(121)
+			_im = _ax.imshow(np.abs(_E_out_ip), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()])
+			_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='Amplitude')
+			_ax2 = _fig.add_subplot(122)
+			_im2 = _ax2.imshow(np.angle(_E_out_ip), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()])
+			_ax2.set_xlim([_xip.min(), _xip.max()])
+			_ax2.set_ylim([_yip.min(), _yip.max()])
+			_ax2.set_xlabel('x [mm]')
+			_ax2.set_ylabel('y [mm]')
+			_fig.colorbar(_im2, ax=_ax2, label='Phase')
+			plt.tight_layout()
+			if save:
+				plt.savefig(save, bbox_inches='tight', transparent=True)
+			plt.show()
+
+		_x, _y = np.meshgrid(_xip, _yip)
+		return np.array([_E_out_ip, _x.real, _y.real])
+
+
+	def freq_array(self, A, dpix):
+		n, m = np.shape(A)
+		
+		fy, fx = np.mgrid[0:m, 0:n]
+
+		fx = fx.flatten()
+		fy = fy.flatten()
+		
+		for i in range(len(fx)):
+			fx[i] = (fx[i] - n*bool(fx[i]>(n-1)/2) ) / (n*dpix)
+			fy[i] = (fy[i] - m*bool(fy[i]>(m-1)/2) ) / (m*dpix)
+		
+		fx = fx.reshape(n,m)
+		fy = fy.reshape(n,m)
+		
+		return fx, fy
+				
+
+	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 minimize(self, a, dz, x, y, og_data):
+		'''
+			Residual function used for fitting the gaussian beam to the input distribution.
+			Input:
+			a 		- 	[float] Guess for w0
+			dz  	- 	[float] Propagation distance, should be known
+			x,y 	- 	[array] meshgrid of the x and y coordinates.
+			og_data - 	[array] input distribution to be fitted.
+		'''
+		return np.abs((self.gaussian_beam_fit(a,dz,x,y) - og_data).flatten()).real
+
+
+	def gaussian_beam_fit(self, a, dz, x, y):
+		'''
+			Function used to fit the Gaussian beam to an input plane. Very similar
+			to the GaussianBeam class below, but modified to allow a best fit.
+			Input:
+			a 	- 	[float] Guess for w0
+			dz  - 	[float] Propagation distance, should be known
+			x,y - 	[array] meshgrid of the x and y coordinates.
+		'''
+		_w0 = a
+		_z 	= dz
+		_r 	  	= np.sqrt(x**2 + y**2)
+		if _z != 0:
+			_z_c 	= np.pi * _w0**2 / self.l
+			_w 		= _w0 * np.sqrt(1 + (_z/_z_c)**2)
+			_R 		= _z + _z_c **2 / _z
+			_phi0 	= np.arctan(_z/_z_c)
+			_k 		= 2 * np.pi / self.l
+			_results = np.sqrt(2 / (np.pi*_w**2)) * np.exp(- _r**2 / _w**2 - 1j * _k * _z - 1j * np.pi * _r**2 / (self.l * _R) + 1j*_phi0)
+		elif _z == 0:
+			_results = np.sqrt(2 / (np.pi*_w0**2)) * np.exp( - _r**2 / _w0**2 )
+		return _results
+
+
+	def optical_transfer_function(self, PSF, x, y, dpix=None):
+		'''
+			Calculates the optical transfer function of the system response
+			Input is the image on the detector, preferably of a PSF
+			The frequencies are given over the whole provided range, but are 
+			typically limited to 1 / (lambda * F/#).
+		'''
+		if dpix is None:
+			dpix = np.mean(np.diff(x.real))
+			if dpix < 1e-99:
+				dpix = np.mean(np.diff((x.real).T))
+		print('dpix is {}'.format(dpix))
+		# FFT input
+		_A_in 			= pyfftw.empty_aligned((PSF.shape), dtype='complex128')
+		_A_in[:]		= PSF
+		
+		# Output of the FFT
+		_A_out 			= pyfftw.empty_aligned(PSF.shape, dtype='complex128')
+		_FFTW_Obj		= pyfftw.FFTW( np.fft.ifftshift(_A_in), _A_out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
+		_FFTW_Obj()
+		_A_out 			= np.fft.fftshift(_A_out)
+
+		# Normalise the output
+		_OTF = _A_out / np.max(_A_out)
+		_fx = np.fft.fftshift(np.fft.fftfreq(np.shape(x)[0], d=dpix))
+		_fy = np.fft.fftshift(np.fft.fftfreq(np.shape(y)[0], d=dpix))
+		return _OTF, _fx, _fy
+
+
+	def modulation_transfer_function(self, *args):
+		'''
+			Calculates the modulation transfer function of the system response
+			Input is the image on the detector, preferably of a PSF
+			The frequencies are given over the whole provided range, but are 
+			typically limited to 1 / (lambda * F/#).
+		'''
+		_OTF, _fx, _fy = self.optical_transfer_function(*args)
+		return np.abs(_OTF), _fx, _fy
+
+
+	def encircled_energy_function(self, PSF, x, y, x0=0, y0=0, dpix=None, originalPSF=None):
+		'''
+			Function that returns the encircled energy as a function of distance
+			from the centroid position, defined by x0 and y0.
+
+			Parameters:
+			PSF 	- 	Intensity map of the PSF
+			x,y 	- 	x and y coordinates corresponding to the PSF
+			x0,y0 	- 	Optional; location of the centroid. If not given,
+						x0=0 and y0=0 are assumed.
+			dpix 	- 	Optional; pixel scale, default=1
+
+			Returns:
+			_EEF 	- 	Array with encircled energy values
+			_r 		- 	Array with radius values
+
+		'''
+		if dpix is None:
+			dpix = np.mean(np.diff(x))
+			if dpix < 1e-99:
+				dpix = np.mean(np.diff(x.T))
+
+		if originalPSF is None:
+			_total_energy 	= np.sum(PSF)
+		else:
+			_total_energy 	= np.sum(originalPSF)
+		_r 				= np.sqrt((x.real-x0)**2 + (y.real-y0)**2)
+		_r_ls 			= np.linspace(0,_r.max(),len(_r))
+		_mask 			= [ _r <= i for i in _r_ls]
+		_EEF 			= np.asarray([np.sum(PSF[m]) for m in _mask])/_total_energy
+		_r 				= _r_ls * dpix
+		return _EEF, _r
+
+	@staticmethod
+	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 select_subsection(self, x, y , outputCoordinates):
+		select_1 = np.max(self.find_nearest(x, outputCoordinates[0]))
+		select_2 = np.max(self.find_nearest(x, outputCoordinates[1]))
+		select_3 = np.max(self.find_nearest(y, outputCoordinates[2]))
+		select_4 = np.max(self.find_nearest(y, outputCoordinates[3]))
+		return [select_1, select_2, select_3, select_4]

fourierOptics/gaussian.py

@@ -0,0 +1,113 @@
+import numpy as np
+
+class GaussianBeam:
+	def __init__(self, w0, l):
+		# Confocal distance
+		self.z_c 	 = np.pi * w0**2 / l
+
+		# Far-field divergence angle
+		self.theta_0 = l / (np.pi * w0)
+		self.l 		 = l
+		self.w0 	 = w0
+		self.fwhm 	 = 1.18 * self.theta_0
+
+
+	def fundamental_mode(self, z, n, m, rmax):
+		y, x = np.mgrid[-rmax:rmax:n*1j,-rmax:rmax:n*1j]
+		r 	 = np.sqrt(x**2 + y**2)
+		if z != 0:
+			_w 		= self.beam_radius(z)
+			_R 		= self.curvature(z)
+			_phi0 	= self.phase_shift(z)
+			_k 		= 2 * np.pi / self.l
+			return np.sqrt(2 / (np.pi*_w**2)) * np.exp(- r**2 / _w**2 - 1j * _k * z - 1j * np.pi * r**2 / (self.l * _R) + 1j*_phi0)
+		elif z == 0:
+			return np.sqrt(2 / (np.pi*self.w0**2)) * np.exp( - r**2 / self.w0**2 ) 
+
+
+	def curvature(self, z):
+		return z + self.z_c **2 / z
+
+
+	def beam_radius(self, z):
+		return self.w0 * np.sqrt(1 + (z/self.z_c)**2)
+
+
+	def phase_shift(self, z):
+		return np.arctan(z/self.z_c)
+
+
+	def relative_power(self, r, z):
+		return np.exp(-2 * (r / self.beam_radius(z))**2)
+
+
+
+
+
+class GaussianBeamPropagation(GaussianBeam):
+	def __init__(self, gaussianBeamObject):
+		self.inputBeam 	= gaussianBeamObject
+		self.l 			= self.inputBeam.l
+		self.z_c 		= self.inputBeam.z_c
+		self.theta_0 	= self.inputBeam.theta_0
+		self.w0 		= self.inputBeam.w0
+		self.fwhm 		= self.inputBeam.fwhm
+		self.beam_width = self.w0
+
+
+	def propagate_with_a_lens(self, ABCD, d_out, n, m, d_in=0, rmax=None):
+		'''
+			Propagates the gaussian beam through a lens with the ray transfer matrix ABCD.
+			This follows the equations of Ch. 3.3 of Goldsmith's Quasi-optical Systems
+
+			Parameters:
+			ABCD 		- 	Ray transfer matrix, use ABCD_paraxial_lens to calculate
+							this matrix for a paraxial lens, with focal length f.
+			d_out 		- 	Distance after the lens over which the beam is propagated.
+			d_in 		- 	Distance from the beam waist (w0) to the lens. Default is 0.
+			rmax 		- 	Maximum radius for which to calculate the beam profile.
+							If None, it will be set to three times the beam width at
+							the output distance.
+
+			Output:
+			The beam information. Can be plotted using imshow()
+			The beam width at the given point can be retrieved by using self.beam_width
+
+		'''
+		_A,_B,_C,_D 	= ABCD
+		_focalDistance 	= - ( (_A * d_in + _B) * (_C * d_in + _D) + _A * _C * self.z_c**2 ) / ( (_C * d_in + _D)**2 + _C**2 * self.z_c**2 )
+		_z = d_out - _focalDistance
+
+		if d_out > _focalDistance:
+			d_l 		= d_out - _focalDistance
+			d_out 		= _focalDistance
+			_w0_out 	= self.w0 / np.sqrt((_C * d_in + _D)**2 + _C**2 * self.z_c**2)
+			_z_c_new 	= np.pi * _w0_out**2 / self.l
+			_wl_out 	= _w0_out * np.sqrt(1 + (d_l/_z_c_new)**2)
+
+		# _q_out 			= ((_A + _C * d_out )*1j*self.z_c + (_A + _C * d_out)*d_in + (_B + _D*d_out)) / (_C * 1j * self.z_c + _C * d_in + _D)
+		_w0_out 		= self.w0 / np.sqrt((_C * d_in + _D)**2 + _C**2 * self.z_c**2)
+
+		self.z_c 		= np.pi * _w0_out**2 / self.l
+		self.w0 		= _w0_out
+		self.theta_0 	= self.l / (np.pi * _w0_out)
+		self.fwhm 		= 1.18 * self.theta_0
+		self.beam_width = self.beam_radius(_z)
+
+		if rmax is None:
+			rmax = 3 * self.beam_width
+
+		_y, _x 			= np.mgrid[-rmax:rmax:n*1j,-rmax:rmax:n*1j]
+		_r 	 			= np.sqrt(_x**2 + _y**2)
+		if _z != 0:
+			_w 		= self.beam_radius(_z)
+			_R 		= self.curvature(_z)
+			_phi0 	= self.phase_shift(_z)
+			_k 		= 2 * np.pi / self.l
+			return np.sqrt(2 / (np.pi*_w**2)) * np.exp(- _r**2 / _w**2 - 1j * _k * _z - 1j * np.pi * _r**2 / (self.l * _R) + 1j*_phi0)
+		elif _z == 0:
+			return np.sqrt(2 / (np.pi*self.w0**2)) * np.exp( - r**2 / self.w0**2 ) 
+
+
+	def ABCD_paraxial_lens(self, f):
+		return [1, 0, -1/f, 1]

fourierOptics/misc.py

@@ -0,0 +1,95 @@
+import numpy as np
+import pickle
+
+'''
+	Miscellaneous functions that do not belong specifically to any of the 
+	other files.
+
+	Author: 	Joost van den Born
+	Contact: 	born@astron.nl
+'''
+
+
+def perform_zero_padding(A, xs, ys):
+	'''
+		Add zero padding around the array
+
+		Input:
+		A 		- 	2D-Input array
+		xs,ys 	- The new shape after padding.
+
+		Output:
+				- Array with the new padded zeros
+
+	'''
+	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):
+	'''
+		For a value it finds the nearest element in an array
+		Returns the index of that element in the original 
+		array.
+	'''
+	array = np.asarray(array)
+	diff = (np.abs(array - value)).argmin()
+	idx = np.unravel_index(diff, array.shape)
+	return idx
+
+def find_maximum(array):
+	'''
+		Finds the index of the maximum of the given array
+	'''
+	idx = np.unravel_index(array.argmax(), array.shape)
+	return idx
+
+def select_subsection(x, y , outputCoordinates):
+	'''
+		Returns the indices of a subsection of an x,y meshgrid
+	'''
+	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):
+	'''
+		Performs linear extrapolation of a meshgrid to a new 
+		shape n,m
+
+		Returns the new extrapolated x,y meshgrid
+	'''
+
+	# 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 save_variable(filename, variable):
+	'''
+		Dumps the given variable to a binary pickle file
+		for later use.
+	'''
+	with open(filename, 'wb') as f:
+		pickle.dump(variable, f, protocol=-1 )
+
+
+def load_variable(filename):
+	'''
+		Loads a binary pickle file for use.
+	'''
+	with open(filename, 'rb') as f:
+		variable = pickle.load(f)
+	return variable
+
+

fourierOptics/opticalTrain.py

@@ -0,0 +1,103 @@
+import numpy as np
+from tabulate import tabulate
+
+
+class OpticalTrain:
+	'''
+		The main class where an optical system is build up sequentially.
+	'''
+	def __init__(self):
+		self.opticalTrain 		= []
+		self.sequentialDistance = []
+
+
+	def add_object(self, obj, dist):
+		'''
+			Adding a lens or mirror object to the opticalTrain
+			obj 	The lens or mirror object
+			dist 	The distance from the the previous object.
+		'''
+		assert type(dist) == float or type(dist) == int, 'Distance is not a float or integer.'
+		self.opticalTrain.append(obj)
+		self.sequentialDistance.append(float(dist))
+
+
+	def print_lens_data(self):
+		'''
+			Print the data of the current opticalTrain.
+		'''
+		_tableData = []
+		for _lens in self.opticalTrain:
+			_i = self.opticalTrain.index(_lens)
+			_distances = self.sequentialDistance
+			_distances.append('-')
+			_tableData.append([_i, _lens.name, _lens.type, _distances[_i+1] ])
+		print(tabulate(_tableData, headers=['#', 'Name', 'Type', 'Distance to next surface']))
+
+
+	def fresnel_number(self, a, z, l):
+		return a**2 / (z * l)
+
+
+	def propagate_system(self, U_in, x, y, l, output=None, paraxial=True):
+		# We need to have an input distribution
+		# We need to determine what propagation method is required
+		# Call the objects in the self.opticalTrain sequentially
+		# Call the subsequent phaseDelay
+		if output == 'complete':
+			_wf_list 	= []
+			_x_list 	= []
+			_y_list 	= []
+			_wf_list.append(_wf)
+			_x_list.append(x)
+			_y_list.append(y)
+		
+		_fou = FourierOptics(l=l)
+
+		_wf = U_in
+
+		oC = [x.min(), x.max(), y.min(), y.max()]
+		show = False
+		_wf, x, y = _fou.fresnel_propagation(_wf, x, y, dz=self.sequentialDistance[0], oversample=None, show=show, save=False, outputCoordinates=oC)
+		if output == 'complete':
+			_wf_list.append(_wf)
+			_x_list.append(x)
+			_y_list.append(y)
+		
+		
+		for _lens in self.opticalTrain[:-1]: # Except for image plane
+
+			title = f"{self.opticalTrain[:-1].index(_lens)+1:02d}"
+			print (' Calculating lens {0} of {1}'.format(title, f"{len(self.opticalTrain[:-1]):02d}"), end="\r")
+
+			_i 	= self.opticalTrain.index(_lens)
+			_dz = self.sequentialDistance[_i+1]
+			_wf *= _lens.phase_delay(x, y, l, paraxial=paraxial)
+
+			# Perform Fresnel number test:
+			_F = self.fresnel_number(_lens.CA, _dz, l)
+			_F = 1
+
+			
+			if _F < 5:
+				# Fresnel propagation
+				if _lens == self.opticalTrain[-2]:
+					oC = [-2, 2, -2, 2]
+					show=True
+				_wf, x, y = _fou.fresnel_propagation(_wf, x, y, dz=self.sequentialDistance[_i+1], oversample=None, show=show, save=False, outputCoordinates=oC)
+			elif _F >= 5:
+				print('Using Plane Wave!')
+				# Angular spectrum (p2p) propagation
+				# _wf, x, y = self.plane_wave_propagation()
+
+			pass
+			if output == 'complete':
+				_wf_list.append(_wf)
+				_x_list.append(x)
+				_y_list.append(y)
+		
+		if output == 'Complete':
+			return np.array([_wf_list, _x_list, _y_list])
+
+		else:
+			return np.array([_wf, x, y])

fourierOptics/optics.py

@@ -0,0 +1,773 @@
+import numpy as np
+from scipy.optimize import fsolve
+from .data import glassCatalog
+
+class LensObject:
+	def __init__(self, CA, name=None):
+		if name:
+			assert type(name)==str, 'No string for a name.'
+			self.name 	= name
+		else:
+			self.name 	= str(self)
+		self.CA 		= CA
+		self.type 		= None
+		self.f 			= None
+		self.sag 		= None
+		self.wavefrontError = False
+
+	def add_wavefront_error(self, z):
+		assert len(z.shape)==2, 'Wavefront data is not 2D. Cannot continue'
+		self.sag = z
+		self.wavefrontError = True
+
+
+
+
+
+class ParaxialLens(LensObject):
+	def __init__(self, f, CA, name=None):
+		super().__init__(CA, name)
+		self.f 			= f
+		self.type 		= 'PARAXIAL'
+
+
+	def find_focal_length(self):
+		return self.f
+
+
+	def phase_delay(self, x, y, l, **kwargs):
+		_k = 2 * np.pi / l
+		_t_l = np.exp(1j * _k / (2 * self.f) * (x**2 + y**2))
+
+		if self.wavefrontError:
+			_t_l *=  np.exp(-1j * _k * self.sag / l)
+
+		# 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)
+		# _t_l *= _lensMask
+		return _t_l
+
+
+
+
+
+class SphericalLens(LensObject):
+	'''
+		The SphericalLens object describes any spherical lens and 
+		finds the phase delay for all light traveling through it.
+		The refractive index, of course, depends on the wavelength.
+	'''
+	def __init__(self, glass, thickness, R1, R2, CA, name=None):
+		super().__init__(CA, name)
+		assert glassCatalog.catalog.get(glass), 'Glass type not in catalog'
+		self.type 		= 'SPHERICAL'
+		self.thickness 	= thickness			# mm
+		self.R1 		= R1 				# mm
+		self.R2 		= R2 				# mm
+		self.glass 		= glass
+		self.dispersionCoefficients = glassCatalog.catalog.get(self.glass)
+		self.referenceWavelength = 1.35e-3 	# mm
+		self.f 			= self.find_focal_length()
+
+
+	def set_first_surface_sag(self, x,y):
+		'''
+			Defines the shape of the first surface and bound the corresponding sag function to
+			the lens object.
+
+			Input:
+			x,y 	- 	Coordinates of the sag.
+		'''
+		# A spherical lens surface requires only the radius of curvature
+		_first_surface_sag 	= self.R1 * (1 - np.sqrt(1 - (x**2 + y**2) / self.R1**2 ) )
+		dist_from_center 	= np.sqrt(x**2 + y**2)
+		_lensMask 			= (dist_from_center <= (self.CA/2)).astype(int)
+		_first_surface_sag *= _lensMask	
+		self.first_surface_sag = _first_surface_sag
+		return self.first_surface_sag
+
+
+	def set_second_surface_sag(self, x,y):
+		'''
+			Defines the shape of the second surface and bound the corresponding sag function to
+			the lens object.
+
+			Input:
+			x,y 	- 	Coordinates of the sag.
+		'''
+		# A spherical lens surface requires only the radius of curvature
+		_second_surface_sag 	= self.R2 * (1 - np.sqrt(1 - (x**2 + y**2) / self.R2**2 ) )
+		dist_from_center 	= np.sqrt(x**2 + y**2)
+		_lensMask 			= (dist_from_center <= (self.CA/2)).astype(int)
+		_second_surface_sag *= _lensMask	
+		self.second_surface_sag = _second_surface_sag
+		return self.second_surface_sag
+
+
+	def get_first_surface_sag(self):
+		return self.first_surface_sag
+
+
+	def get_second_surface_sag(self):
+		return self.second_surface_sag
+
+
+	def find_refractive_index(self, l):
+		# l in mm needs to be converted to l in micron
+		l *= 1e3
+		_a = self.dispersionCoefficients
+		_n = np.sqrt( (_a[0]* l**2)/(l**2 - _a[1]) + (_a[2] * l**2) / (l**2 - _a[3]) + (_a[4] * l**2) / (l**2 - _a[5]) + 1)
+		return _n
+
+
+	def find_focal_length(self):
+		_n 		= self.find_refractive_index(self.referenceWavelength)
+		_f_recip = (_n - 1) * (1 / self.R1 - 1 / self.R2)
+		return 1/_f_recip
+
+
+	def phase_delay(self, x, y, l, paraxial=False):
+		'''
+			This function calculates the phase delay introduced by the lens.
+
+			Input:
+			x, y 		- 	Coordinates of the lens
+			l 			- 	Wavelength in millimeters
+			paraxial 	-	[bool] Apply a paraxial approximation?
+		'''
+		_k = 2 * np.pi / l
+		_n = self.find_refractive_index(l)
+
+		if not paraxial:
+			# For all non-paraxial lenses:
+			# Retrieve the sag of the lens surfaces for coordinates x,y
+			_A1 = self.set_first_surface_sag(x,y)
+			_A2 = self.set_second_surface_sag(x,y)
+			_dxy = self.thickness - _A1 + _A2
+			_t_l = np.exp(-1j * _k * self.thickness) * np.exp(-1j * _k * (_n - 1) * _dxy)
+		elif paraxial:
+			_t_l = np.exp(-1j * _k * self.thickness) * np.exp(1j * _k / (2 * self.f) * (x**2 + y**2))
+
+		if self.wavefrontError:
+			_t_l *=  np.exp(-1j * _k * (_n - 1) * self.sag)
+
+		# 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)
+		_t_l *= _lensMask
+		return _t_l
+
+
+
+
+
+class EvenAsphereLens(LensObject):
+	'''
+		The EvenAsphereLens object describes an even aspherical lens and 
+		finds the phase delay for all light traveling through it.
+		The refractive index, of course, depends on the wavelength.
+	'''
+	def __init__(self, glass, thickness, R1, R2, k1, k2, a1, a2, CA, name=None):
+		super().__init__(CA, name)
+		assert glassCatalog.catalog.get(glass), 'Glass type not in catalog'
+		self.type 			 			= 'EVEN ASPHERE'
+		self.thickness 		 			= thickness			# mm
+		self.R1 			 			= R1 				# mm
+		self.R2 			 			= R2 				# mm
+		self.conicConstantS1 			= k1
+		self.conicConstantS2 			= k2
+		self.asphereCoefficientS1 		= a1
+		self.asphereCoefficientS2 		= a2
+		self.glass 			 			= glass
+		self.dispersionCoefficients 	= glassCatalog.catalog.get(self.glass)
+		self.referenceWavelength 		= 1.35 	# micron
+		self.f 							= self.find_focal_length()
+
+
+	@staticmethod
+	def calculate_surface_sag(x, y, R, k, a):
+		'''
+			Defines the sag shape of the lens
+			Follows the ZEMAX definition on Even Aspheres
+
+			Input parameters:
+			x, y 	-	x,y coordinates of the lens in mm
+			R 		- 	Radius of curvature of the surface in mm
+			k 		- 	Conic constant of the surface
+			a 		- 	Coefficients of the polynomial expansion of the lens sag.
+		'''
+		_r = np.sqrt(x**2 + y**2)
+		_c = 1 / R
+		_expansion = np.sum(np.asarray([a[i] * _r**( 2*(i+1) ) for i in range(len(a))]))
+		_z = _c * _r**2 / ( 1 + np.sqrt( 1 - (1 + k) * _c**2 * _r**2) ) + _expansion
+		return _z
+
+
+	def set_first_surface_sag(self, x,y):
+		'''
+			Defines the shape of the first surface and bound the corresponding sag function to
+			the lens object.
+
+			Input:
+			x,y 	- 	Coordinates of the sag.
+		'''
+		# A spherical lens surface requires only the radius of curvature
+		_first_surface_sag 	= self.calculate_surface_sag(x, y, self.R1, self.conicConstantS1, self.asphereCoefficientS1)
+		dist_from_center 	= np.sqrt(x**2 + y**2)
+		_lensMask 			= (dist_from_center <= (self.CA/2)).astype(int)
+		_first_surface_sag *= _lensMask	
+		self.first_surface_sag = _first_surface_sag
+		return self.first_surface_sag
+
+
+	def set_second_surface_sag(self, x,y):
+		'''
+			Defines the shape of the second surface and bound the corresponding sag function to
+			the lens object.
+
+			Input:
+			x,y 	- 	Coordinates of the sag.
+		'''
+		# A spherical lens surface requires only the radius of curvature
+		_second_surface_sag	= self.calculate_surface_sag(x, y, self.R2, self.conicConstantS2, self.asphereCoefficientS2)
+		dist_from_center 	= np.sqrt(x**2 + y**2)
+		_lensMask 			= (dist_from_center <= (self.CA/2)).astype(int)
+		_second_surface_sag *= _lensMask	
+		self.second_surface_sag = _second_surface_sag
+		return self.second_surface_sag
+
+
+	def get_first_surface_sag(self):
+		return self.first_surface_sag
+
+
+	def get_second_surface_sag(self):
+		return self.second_surface_sag
+
+
+	def find_refractive_index(self, l):
+		# l in mm needs to be converted to l in micron
+		l *= 1e3
+		_a = self.dispersionCoefficients
+		_n = np.sqrt( (_a[0]* l**2)/(l**2 - _a[1]) + (_a[2] * l**2) / (l**2 - _a[3]) + (_a[4] * l**2) / (l**2 - _a[5]) + 1)
+		return _n
+
+
+	def find_focal_length(self):
+		'''
+			Calculates the focal length, where we assume that the 
+			asphere component is small, so that the focal length is 
+			approximately the same as for the spherical lens with
+			radii of curvature of R1 and R2.
+		'''
+		_n 		= self.find_refractive_index(self.referenceWavelength)
+		_f_recip = (_n - 1) * (1 / self.R1 - 1 / self.R2)
+		return 1/_f_recip
+
+
+	def phase_delay(self, x, y, l, paraxial=False):
+		'''
+			This function calculates the phase delay introduced by the lens.
+
+			Input:
+			x, y 		- 	Coordinates of the lens
+			l 			- 	Wavelength in millimeters
+			paraxial 	-	[bool] Apply a paraxial approximation?
+		'''
+		_k = 2 * np.pi / l
+		_n = self.find_refractive_index(l)
+
+		if not paraxial:
+			# For all non-paraxial lenses:
+			# Retrieve the sag of the lens surfaces for coordinates x,y
+			_A1 = self.set_first_surface_sag(x,y)
+			_A2 = self.set_second_surface_sag(x,y)
+			_dxy = self.thickness - _A1 + _A2
+			_t_l = np.exp(-1j * _k * self.thickness) * np.exp(-1j * _k * (_n - 1) * _dxy)
+		elif paraxial:
+			_t_l = np.exp(-1j * _k * self.thickness) * np.exp(1j * _k / (2 * self.f) * (x**2 + y**2))
+
+		if self.wavefrontError:
+			_t_l *=  np.exp(-1j * _k * (_n - 1) * self.sag)
+
+		# 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)
+		_t_l *= _lensMask
+		return _t_l
+
+
+
+
+
+class MirrorObject:
+	'''
+		Describes any mirror. No wavelength dependence exist. 
+		Coatings are not considered.
+	'''
+	def __init__(self, dx, dy, name=None):
+		self.dx 			= dx
+		self.dy 			= dy
+		if name:
+			assert type(name)==str, 'No string for a name.'
+			self.name 		= name
+		else:
+			self.name 		= str(self)
+		self.type 			= None
+		self.sag 			= None
+		self.wavefrontError = False
+
+
+	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
+
+
+
+
+
+class ParabolicMirror(MirrorObject):
+	def __init__(self, f, CA, dx=None, dy=None, name=None):
+		MirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.f = f
+		self.CA = CA
+		self.type = 'PARABOLIC MIRROR'
+
+
+	def phase_delay(self, x, y, l, **kwargs):
+		_k 		= 2 * np.pi / l
+		_dz 	= (x**2 + y**2) / (4 * self.f)
+		
+		if self.wavefrontError:
+			_sagPhase = self.sag
+			_t_m = np.exp(1j * _k * 2 * (_dz + _sagPhase) )
+		else:
+			_t_m = np.exp(1j * _k * 2 * _dz)
+
+		# 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)
+		_t_m *= _mirrorMask
+		return _t_m
+
+
+
+
+
+class SphericalMirror(MirrorObject):
+	def __init__(self, R, CA, dx=None, dy=None, name=None):
+		MirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.R 		= R
+		self.f 		= R/2
+		self.CA 	= CA
+		self.type 	= 'SPHERICAL MIRROR'
+
+
+	def phase_delay(self, x, y, l, **kwargs):
+		_k = 2 * np.pi / l
+		_dz = self.R * (1 - np.sqrt(1 - (x**2 + y**2)/self.R**2) )
+		
+		if self.wavefrontError:
+			_sagPhase = self.sag
+			_t_m = np.exp(1j * _k * 2 * (_dz + _sagPhase) )
+		else:
+			_t_m = np.exp(1j * _k * 2 * _dz)
+
+		# 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)
+		_t_m *= _mirrorMask
+		return _t_m
+
+
+
+
+
+class FreeformMirror(MirrorObject):
+	def __init__(self, x, y, z, dx=None, dy=None, name=None):
+		MirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.x = x
+		self.y = y
+		self.z = z
+		self.type = 'USER DEFINED MIRROR'
+
+
+	def phase_delay(self, l, **kwargs):
+		_k = 2 * np.pi / l
+		if self.wavefrontError:
+			_t_m = np.exp(1j * _k * 2 * ( self.z + self.sag) )
+		else:
+			_t_m = np.exp(1j * _k * 2 * self.z)
+		return _t_m
+
+
+
+
+
+class OpticalSurface:
+	def __init__(self, mask=None, name=None):
+		self.type = 'Dummy surface'
+		if mask:
+			self.x, self.y, self.mask = mask
+			self.type = 'Mask'
+		if name:
+			assert type(name)==str, 'No string for a name.'
+			self.name = name
+
+	def phase_delay(self, x, y, l, **kwargs):
+		return np.ones(np.shape(x))
+
+
+
+
+
+class ADCObject:
+	def __init__(self, wedge=(1.6554, 6.28927, -1.48306)):
+		self.p1 = Prism('S-LAH71', 15, wedge[0], wedge[1], CA=125, name='MCD-ADC-P1')
+		self.p2 = Prism('S-FPL51', 15, wedge[1], wedge[2], CA=125, name='MCD-ADC-P2')
+		self.p3 = Prism('S-FPL51', 15, wedge[2], wedge[1], CA=125, name='MCD-ADC-P3')
+		self.p4 = Prism('S-LAH71', 15, wedge[1], wedge[0], CA=125, name='MCD-ADC-P4')
+		self.theta = None
+		self.set_orientation(90)
+
+
+	def set_orientation(self, theta):
+		self.p1.set_prism_orientation(theta)
+		self.p2.set_prism_orientation(theta)
+		self.p3.set_prism_orientation(180-1*theta)
+		self.p4.set_prism_orientation(180-1*theta)
+		self.theta = theta
+
+
+	def set_field_rotation(self, FR):
+		self.p1.set_prism_orientation(self.theta-FR)
+		self.p2.set_prism_orientation(self.theta-FR)
+		self.p3.set_prism_orientation(180-1*self.theta-FR)
+		self.p4.set_prism_orientation(180-1*self.theta-FR)
+
+
+	def get_phase_delay(self, x, y, l, theta=None, paraxial=False, prism=None):
+		if theta is not None:
+			self.set_orientation(theta)
+		else:
+			theta = self.theta
+		if not prism:
+			_t_p1 = self.p1.phase_delay(x, y, l)#, theta=theta, paraxial=paraxial)
+			_t_p2 = self.p2.phase_delay(x, y, l)#, theta=theta, paraxial=paraxial)
+			_t_p3 = self.p3.phase_delay(x, y, l)#, theta=180-theta, paraxial=paraxial)
+			_t_p4 = self.p4.phase_delay(x, y, l)#, theta=180-theta, paraxial=paraxial)
+			return _t_p1, _t_p2, _t_p3, _t_p4
+		elif prism == 1:
+			_t_p1 = self.p1.phase_delay(x, y, l)#, theta=theta, paraxial=paraxial)
+			return _t_p1
+		elif prism == 2:
+			_t_p2 = self.p2.phase_delay(x, y, l)#, theta=theta, paraxial=paraxial)
+			return _t_p2
+		elif prism == 3:
+			_t_p3 = self.p3.phase_delay(x, y, l)#, theta=180-1*theta, paraxial=paraxial)
+			return _t_p3
+		elif prism == 4:
+			_t_p4 = self.p4.phase_delay(x, y, l)#, theta=180-1*theta, paraxial=paraxial)
+			return _t_p4
+		else:
+			print('This prism ({}) does not exist.'.format(prism))
+		
+
+
+	def add_to_optical_train(self, opticalTrainObject, d, paraxial=True):
+		'''
+			Adds the whole ADC object to an optical train object.
+			Because the ADC consists of multiple prisms, where each
+			prism has two surfaces, this is a bit more involved than
+			just doing opticalTrainObjec.add_object(ADC, dist=d)
+
+			Input:
+			opticalTrainObject 	-	optical train object, see class OpticalTrain
+			d 					- 	distance in mm to next object after the ADC
+			paraxial 			- 	[bool] if true, each ADC only acts as a phase
+									transformation. No diffraction will happen
+									inside the prism.
+		'''
+		if not paraxial:
+			print('This feature has not been implemented yet. Please use paraxial mode.')
+			pass
+
+		if paraxial:
+			opticalTrainObject.add_object(self.p1, dist=d)
+			opticalTrainObject.add_object(self.p2, dist=1)
+			opticalTrainObject.add_object(self.p3, dist=8)
+			opticalTrainObject.add_object(self.p4, dist=1)
+
+	def calculate_dispersion(self, l1, l2, hx=0, hy=0, unit='millimeter'):
+		'''
+			Calculates the dispersion for the ADC (using geometric optics analysis)
+			at its current position, for a normalized on sky field angle (hx, hy). 
+			Default field position is in the center of the FoV.
+
+			Input:
+			l1, l2 	-	Wavelength range of interest (in millimeter)
+			hx, hy 	- 	Field of view of the instrument (normalized between -1 and 1)
+
+			Output:
+			dispX, dispY 	- 	The dispersion in sky angles in milli arc seconds.
+
+			This particular ADC model is only valid for an ADC located in the pupil of the instrument.
+			Change the pupil ratio (Diameter_instrument_pupil / Diameter_Entrance_Pupil) accordingly
+
+		'''
+		if unit == 'millimeter':
+			l1 *= 1e3
+			l2 *= 1e3
+		elif unit == 'micron':
+			pass
+
+		_prismList = [self.p1, self.p2, self.p3, self.p4]
+		_d1x_list, _d1y_list = [], []
+		# _pupilRatio = 0.07943/38.542
+		_pupilRatio = 0.0801/38.542
+
+		hx = hx * 7.5e-3
+		hy = hy * 7.5e-3
+		ix = hx * 1/_pupilRatio
+		iy = hy * 1/_pupilRatio
+
+		for l in [l1,l2]:
+			_deltaListX = np.zeros(len(_prismList))
+			_deltaListY = np.zeros(len(_prismList))
+			for p in range(4):
+				_n_g = _prismList[p].find_refractive_index(l, unit='micron')
+				
+				_b1x = _prismList[p].tilt1x
+				_b1y = _prismList[p].tilt1y
+
+				_b2x = _prismList[p].tilt2x
+				_b2y = _prismList[p].tilt2y
+
+				if p==0:
+				    _t_iz_n_x = ix # Zero before first prism
+				    _t_iz_n_y = iy # Zero before first prism
+
+				# X direction    
+				_t_i1_n_x = _t_iz_n_x #+ _b1x
+				_d1x = self.calculate_deviation(_t_i1_n_x, _b1x, _b2x, _n_g)
+				_t_tz_n_x = _t_iz_n_x - _d1x
+				_t_iz_n_x = _t_tz_n_x         # Set output to new input (for the next prism)
+
+				# Y direction
+				_t_i1_n_y = _t_iz_n_y #+ _b1y
+				_d1y = self.calculate_deviation(_t_i1_n_y, _b1y, _b2y, _n_g)
+				_t_tz_n_y = _t_iz_n_y - _d1y
+				_t_iz_n_y = _t_tz_n_y         # Set output to new input (for the next prism)
+
+				# Adding the deviations to the list with delta terms
+				_deltaListX[p] = _d1x 
+				_deltaListY[p] = _d1y
+			_d1x_list.append(_deltaListX.sum())
+			_d1y_list.append(_deltaListY.sum())
+
+		self.dispX = ( _d1x_list[1] - _d1x_list[0] ) * 3600 * 1000 * _pupilRatio
+		self.dispY = ( _d1y_list[1] - _d1y_list[0] ) * 3600 * 1000 * _pupilRatio
+		return self.dispX, self.dispY
+
+
+	@staticmethod
+	def calculate_deviation(t_i1, b1, b2, n):
+		'''
+			Returns deviation of the ray after the prism with apex angle alpha
+			only now using Snell's law only. Using the formula from Hagen & Tkazcyk (2011)
+			Units: degrees
+		'''
+		t_i1 	= np.deg2rad(t_i1)
+		b1 		= np.deg2rad(b1)
+		b2 		= np.deg2rad(b2)
+		alpha 	= b2 - b1
+		delta 	= t_i1 - b2 - np.arcsin( n * np.sin( np.arcsin(1/n * np.sin(t_i1 - b1) ) - alpha ) )
+		return np.degrees(delta)
+
+	def find_optimum_angle(self, inputDispersion, l1, l2, x0=50):
+		'''
+			Simple minimizer to find the (local) minimum for a best position of the ADC.
+		'''
+		_solver = fsolve(self.optimizer, x0=x0, args=(self, inputDispersion, l1, l2))
+		return _solver[0]
+	
+	@staticmethod
+	def optimizer(theta, adcObject, inputDispersion, l1, l2, FR=None):
+		'''
+			Minimization function to find the optimum rotation angle
+			Input:
+			theta 		- 	Angle in degrees relative to initial position
+			adcObject 	- 	the ADC class
+			inputDispersion - 	The amount of dispersion to correct for (in mas)
+			l1, l2 		- 	The wavelength range to correct for (micrometer)
+			FR 			- 	Field rotation, not used currently.
+		'''
+		if FR:
+			FRdeg = np.deg2rad(FR)
+			inputDispersionX = np.cos(FRdeg) * inputDispersion
+			inputDispersionY = np.sin(FRdeg) * inputDispersion
+			adcObject.set_orientation(float(theta))
+			ADCdisp = adcObject.calculate_dispersion(l1,l2)
+			return np.sqrt((ADCdisp[1] - inputDispersionY)**2 + (ADCdisp[0] - inputDispersionX)**2)
+		else:	
+			adcObject.set_orientation(float(theta))
+			ADCdisp = adcObject.calculate_dispersion(l1,l2)
+			ADCdisp = ADCdisp[1]
+			return ADCdisp - inputDispersion
+
+
+
+class Prism:
+	def __init__(self, glass, thickness, tilt1, tilt2, CA, name=None):
+		if name:
+			assert type(name)==str, 'No string for a name.'
+			self.name 	= name
+		else:
+			self.name 	= str(self)
+		self.type 		= 'PRISM'
+		self.glass 		= glass
+		self.thickness 	= thickness
+		self.tilt1x		= tilt1
+		self.tilt1y 	= 0
+		self.tilt2x		= tilt2
+		self.tilt2y 	= 0
+		self.CA 		= CA
+		self.sag_combined = None
+		self.sag1 		= None
+		self.sag2		= None
+		self.wavefrontError = False
+		self.thetaPrism	= 0
+
+		# Remember original tilts for reset:
+		self.tilt1x_OG	= tilt1
+		self.tilt1y_OG 	= 0
+		self.tilt2x_OG	= tilt2
+		self.tilt2y_OG 	= 0
+
+
+	def find_refractive_index(self, l, unit='millimeter'):
+		if unit == 'millimeter':
+			l *= 1e3
+		elif unit == 'micron':
+			pass
+		if self.glass == 'S-FPL51':
+			a = [1.170105050E+00, 4.757107830E-02, 7.638324450E-01]
+			b = [6.162039240E-03, 2.633728760E-02, 1.418826420E+02]
+		elif self.glass == 'S-LAH71':
+			a = [1.982800310E+00, 3.167584500E-01, 2.444726460E+00]
+			b = [1.189874560E-02, 5.271560010E-02, 2.132206970E+02]
+		return self.calculate_refractive_index(l, a, b, unit='micron') 
+
+
+	def calculate_refractive_index(self, l, a, b, unit='millimeter'):
+		# l in mm needs to be converted to l in micron
+		if unit == 'millimeter':
+			l *= 1e3
+		elif unit == 'micron':
+			pass
+		n = np.sqrt( (a[0]* l**2)/(l**2 - b[0]) + (a[1] * l**2) / (l**2 - b[1]) + (a[2] * l**2) / (l**2 - b[2]) + 1)
+		return n
+
+
+	def set_prism_orientation(self, theta):
+		self.tilt1x, self.tilt1y = self.tilt_rotation(self.tilt1x_OG, self.tilt1y_OG, theta)
+		self.tilt2x, self.tilt2y = self.tilt_rotation(self.tilt2x_OG, self.tilt2y_OG, theta)
+		self.thetaPrism = theta
+	
+
+	def tilt_rotation(self, tx,ty, theta):
+		theta = np.deg2rad(theta-90)
+		tx_new = tx * np.cos(theta) - ty * np.sin(theta)
+		ty_new = tx * np.sin(theta) + ty * np.cos(theta)
+		return tx_new, ty_new
+
+
+	def get_surface_sag(self, x, y, theta):
+		# self.set_prism_orientation(theta)
+		self.first_surface_sag = y * np.tan(np.deg2rad(self.tilt1x)) + x * np.tan(np.deg2rad(self.tilt1y))
+		self.second_surface_sag = y * np.tan(np.deg2rad(self.tilt2x)) + x * np.tan(np.deg2rad(self.tilt2y))
+		
+		# Masking out the parts that fall outside the prism clear aperture
+		dist_from_center = np.sqrt(x**2 + y**2)
+		_prismMask = (dist_from_center <= (self.CA/2)).astype(int)
+		self.first_surface_sag 	*= _prismMask
+		self.second_surface_sag *= _prismMask
+		return self.first_surface_sag, -1*self.second_surface_sag
+
+
+	def add_wavefront_error(self, z, surface=0):
+		assert len(z.shape)==2, 'Wavefront data does not have the proper shape. Cannot continue'
+		assert surface in [0,1,2], 'Surface not understood. Use 0 for transmitted wavefront error, \
+		1 for surface 1 or 2 for surface 2.'
+		if surface == 0:
+			self.sag_combined = z
+		elif surface == 1:
+			self.sag1 = z
+		elif surface == 2:
+			self.sag2 = z
+		self.wavefrontError = True
+
+
+	def phase_delay(self, x, y, l, theta=None, paraxial=True):
+		'''
+			This function calculates the phase delay introduced by the lens.
+
+			Input:
+			x, y 		- 	Coordinates of the lens
+			l 			- 	Wavelength in millimeters
+			paraxial 	-	[bool] Apply a paraxial approximation? In the
+							case of a prism, this means that the prism acts
+							only as a phase transformation. Diffraction 
+							inside the prism will not take place.
+		'''
+		if theta is None:
+			theta = self.thetaPrism
+
+		_k = 2 * np.pi / l
+		_n = self.find_refractive_index(l)
+
+		# Retrieve the sag of the lens surfaces for coordinates x,y
+		_A1, _A2 	= self.get_surface_sag(x,y, theta)
+
+		# Masking out the parts that fall outside the lens clear aperture
+		dist_from_center = np.sqrt(x**2 + y**2)
+		_prismMask = (dist_from_center <= (self.CA/2)).astype(int)
+
+		if paraxial:
+			# Calculates the phase transformation of this prism. 
+			# Diffraction inside the prism will not take place.
+			_dxy 		= self.thickness + _A1 + _A2 
+			_t_p 		= np.exp(-1j * _k * self.thickness) * np.exp(-1j * _k * (_n - 1) * _dxy)
+			if self.wavefrontError:
+				if self.sag_combined is not None:
+					_t_p *=  np.exp(-1j * _k * (_n - 1) * self.sag_combined)
+				elif self.sag1 is not None:
+					_t_p *=  np.exp(-1j * _k * (_n - 1) * self.sag1)
+				elif self.sag2 is not None:
+					_t_p *=  np.exp(-1j * _k * (_n - 1) * self.sag2)
+
+			# _t_p *= _prismMask
+			return _t_p
+
+		if not paraxial:
+			# Return two different phase transformations
+			# Option to add surface form error (or wavefront error) 
+			# to each interface.
+			_t_p_s1 	= np.exp(-1j * _k * (_n - 1) * _A1)
+			_t_p_s2 	= np.exp(-1j * _k * (_n - 1) * _A2)
+
+			if self.wavefrontError:
+				if self.sag1 is not None:
+					_t_p_s1 *=  np.exp(-1j * _k * (_n - 1) * self.sag1)
+				if self.sag2 is not None:
+					_t_p_s2 *=  np.exp(-1j * _k * (_n - 1) * self.sag2)
+				if self.sag_combined is not None and self.sag1 is None and self.sag2 is None:
+					_t_p_s1 *=  np.exp(-1j * _k * (_n - 1) * self.sag_combined)
+
+			# _t_p_s1 *= _prismMask
+			# _t_p_s2 *= _prismMask
+			return _t_p_s1, _t_p_s2
+
+		

fourierOptics/profiles.py

@@ -0,0 +1,73 @@
+# from .__init__ import __data_dir__
+import numpy as np
+import os
+from astropy.io import fits
+from scipy import interpolate
+
+
+
+class Profiles:
+	def __init__(self):
+		pass
+
+	def top_hat(self, r, n, m, dpix, oversample=1, dx=0, dy=0):
+		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-dx)**2 + (j-dy)**2)
+		_Anew = (dist_from_center <= (r)).astype(int)
+		return _Anew
+
+
+	def obscured_top_hat(self, r1, r2, n, m, dpix, oversample=1):
+		'''
+			Creates a grid of n x m points, following a top hat function
+			where all points inside a radius r is 1 and 0 otherwise.
+		'''
+		assert r1 > r2, "r2 < r1, change input"
+		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))
+		return _Anew.astype(int)
+
+
+	def spider_top_hat(self, r1, r2, n, m, dpix, oversample=1):
+		'''
+			Creates a grid of n x m points, following a top hat function
+			where all points inside a radius r is 1 and 0 otherwise.
+		'''
+		assert r1 > r2, "r2 < r1, change input"
+		j,i 				= np.mgrid[-n/2*dpix : n/2*dpix : 1j*n*oversample, -m/2*dpix : m/2*dpix : 1j*m*oversample]
+		n *= oversample
+		m *= oversample
+		dist_from_center 	= np.sqrt(i**2 + j**2)
+		_Anew 				= (dist_from_center <= (r1))  & (dist_from_center >= (r2))
+		_Aspider 			= (np.abs(i) < (0.05*r1)) | (np.abs(j) < (0.05*r1))
+		_Anew 				= _Anew & ~_Aspider
+		return _Anew
+
+
+	def gaussian(self, r, n, m, dpix, oversample=1, dx=0, dy=0):
+		'''
+			Creates a grid of n x m points, following a gaussian function
+			where the radius r is the FWHM of the gaussian shape.
+		'''
+		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-dx)**2 + (j-dy)**2)
+		return np.exp(-((j-dx)**2 + (i-dy)**2) / r**2)
+
+
+	def ELT_pupil(self, n, m, dpix):
+		# Find data directory
+		__pkg_dir__ 	= os.path.dirname(os.path.realpath(__file__))
+		__data_dir__ 	= os.path.join(__pkg_dir__, 'data')
+		_fits_image 	= __data_dir__ + '/ELT_M1_pupil_profile.fits'
+		_ELT_pupil 		= fits.open(_fits_image)[0].data
+		_ELT_inner_diam = 38542 # all glass diameter of M1 in mm
+		_dpix_fits 		= _ELT_inner_diam / 623 # pixels on the diameter filled with 'glass'
+		_resx, _resy 	= np.shape(_ELT_pupil)
+		_xf				= np.linspace(-_resx/2*_dpix_fits, _resx/2*_dpix_fits, _resx)
+		_yf 			= np.linspace(-_resy/2*_dpix_fits, _resy/2*_dpix_fits, _resy)
+		_f_pupil		= interpolate.interp2d(_xf, _yf, _ELT_pupil, kind='linear')
+		_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

fourierOptics/psfFitter.py

@@ -0,0 +1,255 @@
+import numpy as np
+from scipy import interpolate
+from scipy.optimize import leastsq
+import types
+import scipy.optimize as opt
+from scipy.special import j1
+
+
+class PSFFitter:
+	'''
+		Wrapper around scipy.optimize.curve_fit(), so that point spread functions can be fitted
+		using some common profiles (e.g. Gaussian, Moffat and Airy pattern profiles).
+	'''
+	def __init__(self, x, y, psfdata):
+		self.x 			= x
+		self.y 			= y
+		self.psfdata 	= psfdata
+
+
+	def gaussian_PSF(self, xy, amplitude, x0, y0, sigma_x, sigma_y, theta, offset):
+		'''
+			Elliptical 2D gaussian profile
+
+			Input:
+			xy 			- 	Combined x,y coordinates. x and y should be a meshgrid
+			amplitude 	- 	amplitude of the profile, i.e. maximum of the profile
+			x0, y0 		- 	shift in x and y of the central peak
+			sigma_x 	-	width of the profile in x
+			sigma_y 	- 	width of the profile in y
+			theta 		- 	Rotation of the elliptical gaussian profile
+			offset 		- 	noise floor or sky background level
+		'''
+		_x,_y 	= xy
+		_x0 	= float(x0)
+		_y0 	= float(y0)    
+		_a 		= (np.cos(theta)**2)/(2*sigma_x**2) + (np.sin(theta)**2)/(2*sigma_y**2)
+		_b 		= -(np.sin(2*theta))/(4*sigma_x**2) + (np.sin(2*theta))/(4*sigma_y**2)
+		_c 		= (np.sin(theta)**2)/(2*sigma_x**2) + (np.cos(theta)**2)/(2*sigma_y**2)
+		_g 		= offset + amplitude*np.exp( - (_a*((_x-_x0)**2) + 2*_b*(_x-_x0)*(_y-_y0) + _c*((_y-_y0)**2)))
+		return _g.ravel()
+
+
+	def fit_gaussian_PSF(self, initialGuess=None):
+		''' 
+			Fit a gaussian PSF profile, a Moffat PSF profile and an Airy pattern
+			An OptimizeWarning is given when the best fit is reached, for some reason; 
+			see also:
+			https://stackoverflow.com/questions/41725377/why-isnt-curve-fit-able-to-estimate-the-covariance-of-the-parameter-if-the-pa
+		
+			Input:
+			initialGuess	- 	Tuple, containing (amplitude, x0, y0, sigma_x, sigma_y, theta, offset)
+		'''
+		_popt, _pcov = opt.curve_fit(self.gaussian_PSF, (self.x,self.y), self.psfdata.flatten(), p0=initialGuess)
+		return self.gaussian_PSF((self.x,self.y), *_popt).reshape(self.x.shape), _popt, _pcov
+
+
+	def moffat_PSF(self, xy, amplitude, x0, y0, fwhm1, fwhm2, theta, offset, beta):
+		'''
+			Elliptical moffat distribution, often used to fit a PSF without diffraction rings
+			Equation taken from http://www.aspylib.com/doc/aspylib_fitting.html
+			As beta goes to infinity, the moffat profile approaches a gaussian; as mentioned
+			here https://www.gnu.org/software/gnuastro/manual/html_node/PSF.html
+
+			Input:
+			xy 			- 	Combined x,y coordinates. x and y should be a meshgrid
+			amplitude 	- 	amplitude of the profile, i.e. maximum of the profile
+			x0, y0 		- 	shift in x and y of the central peak
+			fwhm1, fwhm2 -	width of the profile in perpendicular directions
+			theta 		- 	Rotation of the elliptical moffat profile
+			offset 		- 	noise floor or sky background level
+			beta 		- 	shape parameter of the moffat profile
+		'''
+		_x,_y 	 = xy
+		_x0 	 = float(x0)
+		_y0 	 = float(y0)
+		_alpha_1 = fwhm1 / (2 * np.sqrt(2**(1/beta) - 1))
+		_alpha_2 = fwhm2 / (2 * np.sqrt(2**(1/beta) - 1))
+		_a 		 = (np.cos(theta)/_alpha_1)**2 + (np.sin(theta)/_alpha_2)**2
+		_b 		 = (np.sin(theta)/_alpha_1)**2 + (np.cos(theta)/_alpha_2)**2
+		_c 		 = np.sin(theta) * np.cos(theta) * (1/(_alpha_1**2) - 1/(_alpha_2**2))
+		_m 		 = offset + amplitude/((1 + _a*(_x-_x0)**2 + _b*(_y-_y0)**2 + 2*_c*(_x-_x0)*(_y-_y0))**beta)
+		return _m.ravel()
+
+
+	def fit_moffat_PSF(self, initialGuess=None):
+		'''
+			Fit a Moffat PSF profile
+			As beta goes to infinity, the moffat profile approaches a gaussian; as mentioned
+			here https://www.gnu.org/software/gnuastro/manual/html_node/PSF.html
+			This might raise a RuntimeWarning.
+
+			Input:
+			initialGuess	- 	Tuple, containing (amplitude, x0, y0, fwhm1, fwhm2, theta, offset, beta)
+		'''
+		_popt, _pcov = opt.curve_fit(self.moffat_PSF, (self.x,self.y), self.psfdata.flatten(), p0=initialGuess)
+		return self.moffat_PSF((self.x,self.y), *_popt).reshape(self.x.shape), _popt, _pcov
+
+	
+	def airy_PSF(self, xy, amplitude, x0, y0, offset, radius):
+		'''
+			2D Airy pattern with offset included
+
+			Input:
+			xy 			- 	Combined x,y coordinates. x and y should be a meshgrid
+			amplitude 	- 	amplitude of the profile, i.e. maximum of the profile
+			x0, y0 		- 	shift in x and y of the central peak
+			offset 		- 	noise floor or sky background level
+		'''
+		_x,_y 	= xy
+		_x0 	= float(x0)
+		_y0 	= float(y0)
+		_r 		= np.sqrt((_x-_x0)**2 + (_y-_y0)**2)
+		_Rz 	= 1.2196698912665045
+		_inp 	= (np.pi * _r) / (radius/_Rz)
+		_J1 	= j1(_inp)
+		_fraction = np.nan_to_num(2*_J1 / _inp, nan=1.0)
+		_a 		= offset + amplitude * (_fraction)**2
+		return _a.ravel()
+
+
+	def fit_airy_PSF(self, initialGuess=None):
+		'''
+			Repackaging of scipy's curve_fit to fit an Airy pattern to the data.
+			No arguments are necessary, but an initial guess of the parameters 
+			is recommended.
+
+			Input:
+			initialGuess	- 	Tuple, containing (amplitude, x0, y0, offset, radius)
+		'''
+		_popt, _pcov = opt.curve_fit(self.airy_PSF, (self.x,self.y), self.psfdata.flatten(), p0=initialGuess)
+		return self.airy_PSF((self.x,self.y), *_popt).reshape(self.x.shape), _popt, _pcov
+
+
+	def reference_PSF(self, p, xy, perturbedPSF):
+		'''
+			Residual function used for the mimizer functions in 
+			self.fit_custom_field(), see function below.
+		'''
+		_x0, _y0 	= p
+		_x, _y 	 	= xy
+		_a 		 	= self.f_Eref_real(_x - _x0, _y - _y0) + 1j * self.f_Eref_imag(_x - _x0, _y - _y0)
+		_coupling 	= self.get_coupling(_a, perturbedPSF)
+		return 1 - np.abs(_coupling)**2
+
+
+	def fit_custom_field(self, Eref, Eper=None, method='local', niter_success=10, bounds=None, return_coupling=False):
+		'''
+			Fit any custom PSF to the data (self.psfdata). It requires an input
+			of the custom PSF in the same shape as the data.  Then a global 
+			optimization algorithm (scipy.optimization.differential_evolution)
+			is performed to find the translation parameters; x0 and y0.
+
+			Note:
+			The current implementation is not able to find other parameters, 
+			such as amplitude, amplitude offset or radius. It assumes that 
+			the custom PSF and data are quite similar in shape.
+			Also note that what is here called the PSF is actually the complex
+			field distribution of the PSF, i.e. not the intensity, but the
+			amplitude and phase information of the PSF.
+
+			Input:
+			Eref 		- 	Custom PSF, including phase information
+			Eper 		-	optional, perturbed PSF. By default self.psfdata is taken.
+			method 		- 	Optimization strategy to use. Choose between
+								'local'
+								'global'
+							'local' first applies a local search and is very sensitive
+							to the initial guess.
+							'global' searches a global minimum. This is less sensitive
+							to the initial guess, but requires much more computing
+							time.
+			niter_succes-	From scipy.optimize.basinhopping:
+							Stop the run if the global minimum candidate remains the 
+							same for this number of iterations.
+			bounds 		- 	Bounds for the local minimizer
+
+			Returns:
+			_best_match - 	Best matched PSF amplitude and phase data.
+			_popt 		- 	Optimum parameters, i.e. translation of the centroid.
+			_pvar 		- 	Variances of the parameters.
+			_coupling 	- 	Optional; coupling factor
+
+		'''
+
+		assert Eref.shape == self.psfdata.shape, 'The shape of the new (Eper) and reference PSF (Eref) do not match.'
+		# Set x and y coordinates for interpolation. No 2D meshgrid is necessary.
+		_x = self.x[0]
+		_y = self.y[:,0]
+
+		# Check for custom psf data and custom bounds
+		if Eper is None:
+			Eper = self.psfdata
+			print('taken self.psfdata')
+
+		# Create interpolation function to match amplitude and phase to x and y coordinates.
+		self.f_Eref_real = interpolate.interp2d( _x, _y, Eref.real, kind='cubic')
+		self.f_Eref_imag = interpolate.interp2d( _x, _y, Eref.imag, kind='cubic')
+
+		# Determining a good initial guess
+		# Reasonable to assume >0.9 of the maximum amplitude is near the centroid position.
+		_max_amplitude = np.max(np.abs(Eref))
+		_index_i, _index_j = np.where(np.abs(Eref) > (0.9 * _max_amplitude) )
+		_x0_guess = np.mean(self.x[_index_i, _index_j])
+		_y0_guess = np.mean(self.y[_index_i, _index_j])
+		_p0 = [_x0_guess, _y0_guess]
+
+		if method == 'global':			
+			# Arguments for basinhopping function (i.e. global search function)
+			_min_kwargs = {"args":((_x,_y), Eper), "jac":False}
+			_popt = opt.basinhopping(self.reference_PSF, _p0, minimizer_kwargs=_min_kwargs, niter_success=niter_success)
+			_p0 = _popt.x 		# Set optimal parameters as new first guess.
+			method = 'local' 	# Switch to local minimzer to also calculate the uncertainties
+
+		if method == 'local':
+			# Mimizer function using L-BFGS-B method.
+			if bounds is None:
+				bounds = [(_x.min(), _x.max()),(_y.min(), _y.max())]
+			_res = opt.minimize(self.reference_PSF, _p0, bounds=bounds, args=((_x,_y), Eper) )
+			_popt = _res.x
+
+			# Getting the uncertainty in the minimum
+			_ftol = 2.220446049250313e-09 # ftol assuming L-BFGS-B method in the minimizer.
+			_hess_inv_i = np.diag(_res.hess_inv.todense())
+			_pvar = max(1, abs(_res.fun)) * _ftol * _hess_inv_i
+
+		# Centroid coordinates
+		_x0_opt, _y0_opt = _popt
+		_best_match = self.f_Eref_real(_x - _x0_opt, _y - _y0_opt) + 1j * self.f_Eref_imag(_x - _x0_opt, _y - _y0_opt)
+		
+		if return_coupling:
+			_coupling = self.get_coupling(_best_match, Eper)
+			return _best_match, _popt, _pvar, _coupling
+		else:
+			return _best_match, _popt, _pvar
+
+
+	def get_coupling(self, Eref, Eper):
+		'''
+			Calculates the coupling integral as given by eq. 6.5 in 
+			Quasi-optical systems by Goldsmith.
+
+			This complex value is a measure of the match between a reference
+			field and a perturbed field. If the complex coupling factor is c,
+			then |c|**2 is the degradation of the perturbed field relative to
+			the reference field, not unlike the Strehl ratio.
+
+			Parameters:
+			Eref 	- 	Array; Reference field
+			Eper 	- 	Array; Perturbed field
+
+			Returns:
+			_c 		- 	Complex float; Coupling factor
+		'''
+		_c = np.sum(Eref * np.conj(Eper))/ np.sqrt(np.sum(np.abs(Eref)**2)*np.sum(np.abs(Eper)**2))
+		return _c 

fourierOptics/wavefrontErrors.py

@@ -0,0 +1,506 @@
+import numpy as np
+import pyfftw
+import types
+from math import ceil
+from scipy.special import factorial, j1
+
+
+class WavefrontError:
+	'''
+		This class contains methods to describe wavefront errors;
+		either by means of zernike polynomials (description and fits)
+		or by means of Power Spectral Density analysis.
+	'''
+	def __init__(self, n, m):
+		self.n = n 				# Resolution in x-direction
+		self.m = m 				# Resolution in y-direction
+		self.resolution = (n,m)
+
+
+	def calculate_rnm(self,n,m,rho):
+		'''
+		@author: Tibor Agocs
+		Return an array with the zernike Rnm polynomial calculated at rho points.
+
+
+		**ARGUMENTS:**
+
+			=== ==========================================
+			n    n order of the Zernike polynomial
+			m    m order of the Zernike polynomial
+			rho  Matrix containing the radial coordinates. 
+			=== ==========================================       
+
+		.. note:: For rho>1 the returned value is 0
+		    
+		.. note:: Values for rho<0 are silently returned as rho=0
+
+		'''
+
+		if (n-m)%2!=0:
+			raise Exception('n-m must be even')
+		if abs(m)>n:
+			raise Exception('The following must be true |m|<=n')
+
+		if(n==0 and m==0):
+			data=np.ones(rho.shape)
+			return data
+		_mask		= rho > 1
+		rho[_mask] 	= 0
+		_Rnm 		= np.zeros(rho.shape)
+		_S 			= (n-np.abs(m))/2
+		for s in range (0,int(_S+1)):
+			_CR 	=(-1)**s*factorial(n-s)/ \
+					(factorial(s)*factorial(-s+(n+np.abs(m))/2)* \
+					factorial(-s+(n-np.abs(m))/2))
+			_p 		= _CR * ( rho**(n-2*s) )
+			_Rnm 	= _Rnm + _p
+		_Rnm[_mask] = 0
+		return _Rnm
+
+
+	def zernike_polynomial_polar(self,n,m,rho,theta):
+		'''
+			@author: Tibor Agocs
+			Returns the an array with the Zernike polynomial evaluated in the rho and 
+			theta.
+
+			**ARGUMENTS:** 
+			===== ==========================================     
+			n     n order of the Zernike polynomial
+			m     m order of the Zernike polynomial
+			rho   Matrix containing the radial coordinates. 
+			theta Matrix containing the angular coordinates    
+			===== ==========================================
+
+			.. note:: For rho>1 the returned value is 0
+
+			.. note:: Values for rho<0 are silently returned as rho=0
+
+			**RETURNS**
+			===== ===========================================
+			_Zmn  Array with values of the Zernike polynomial
+			===== ===========================================
+		'''
+
+		_Rnm 	= self.calculate_rnm(n,m,rho)
+
+		# Normalization constant
+		_delta_m0 = bool(m==0)
+		_NC 	= np.sqrt(2*(n+1)/(1+_delta_m0))
+
+		if m >= 0:
+			_Zmn = _NC * _Rnm * np.cos(m*theta)
+		elif m < 0:
+			_Zmn = _NC * _Rnm * np.sin(m*theta)
+		return _Zmn
+
+
+	def zernike_polar(self, coefficient, rho, theta, noll=False, rescale='auto'):
+		'''
+			Returns the full surface created from a list of coefficients,
+			using the Zernike Polynomials.
+
+			rho must be between 0 and 1, if not it will be rescaled,
+			see also zernikePolynomialPolar()
+
+			**ARGUMENTS:** 
+			============ ==========================================     
+			coefficient  n order of the Zernike polynomial
+			rho   		 Matrix containing the radial coordinates. 
+			theta 		 Matrix containing the angular coordinates
+			noll 		 Boolean to indicate to use noll indices    
+			============ ==========================================
+
+			**RETURNS**
+			**RETURNS**
+			===== ===========================================
+			_zP   Array with values of the Zernike polynomial
+			===== ===========================================
+		'''
+		if rescale == 'auto':
+			if np.max(rho) > np.sqrt(2):
+				print('Rescaling!')
+				rho = np.sqrt(2) * rho / np.max(rho)
+
+		_termList = []
+		for i in range(len(coefficient)):
+			_C = coefficient[i]
+			if noll:
+				_jTerm = self.convert_noll_to_standard(i+1)
+			_n,_m = self.i_to_nm(i)
+			_termList.append( _C * self.zernike_polynomial_polar(_n,_m,rho,theta) )
+		_zP = sum(_termList)
+		return _zP
+
+
+	def zernike_cartesian(self, coefficient, x, y, noll=False, addToData=False, **kwargs):
+		'''
+			Converts the cartesian grid x,y into polar
+			coordinates and returns the data, such that
+			it can be plotted with e.g. plt.imshow()
+
+			**ARGUMENTS:** 
+
+			============ ==========================================     
+			coefficient  n order of the Zernike polynomial
+			rho   		 Matrix containing the radial coordinates. 
+			theta 		 Matrix containing the angular coordinates
+			noll 		 Boolean to indicate to use noll indices    
+			============ ==========================================
+
+			**RETURNS**
+			===== ===========================================
+			_zP   Array with values of the Zernike polynomial
+			===== ===========================================
+		'''
+		_x0 	= (np.max(x) - np.min(x)) / 2 + np.min(x)
+		_y0 	= (np.max(y) - np.min(y)) / 2 + np.min(y)
+		_rho   	= np.sqrt(( x - _x0 )**2 + ( y - _y0 )**2 )
+		_theta 	= np.arctan2( ( y - _y0 ) , ( x - _x0 ) )
+		_zP 	= self.zernike_polar(coefficient, _rho, _theta, noll, **kwargs)
+		if not addToData:
+			self.data = _zP	
+		else:
+			self.data += _zP
+		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))
+		return _zP
+
+
+	def i_to_nm(self, i):
+		'''
+			Converts a single index to n and m for the Zernike Polynomials
+		'''
+		n = ceil((-3+np.sqrt(9+8*i))/2)
+		m = 2 * i - n*(n+2)
+		return int(n),int(m)
+
+
+	def convert_noll_to_standard(self, j):
+		'''
+
+		'''
+		conversionList = [1, 3, 2, 5, 4, 6, 9, 7, 8, 10, 15, 13, 11, 12, 14, 21, 19, 17, 16, 18, 20, 27, 25, 23, 22, 24, 26, 28, 35, 33, 31, 29, 30, 32, 34, 36, 45, 43, 41, 39, 37, 38, 40, 42, 44, 55, 53, 51, 49, 47, 46, 48, 50, 52, 54, 65, 63, 61, 59, 57, 56, 58, 60, 62, 64, 66, 77, 75, 73, 71, 69, 67, 68, 70, 72, 74, 76, 78, 91, 89, 87, 85, 83, 81, 79, 80, 82, 84, 86, 88, 90, 105, 103, 101, 99, 97, 95, 93, 92, 94, 96, 98, 100, 102, 104, 119, 117, 115, 113, 111, 109, 107, 106, 108, 110, 112, 114, 116, 118, 120]
+		i = conversionList.index(j)
+		return i
+
+
+	def fit_zernike_coefficients(self, data, x=None, y=None, modes=21):
+		'''
+			Fits the zernike coefficients to a data set, using 
+			that the zernike modes can be treated independently.
+			The modes are set following standard index notation.
+			We solve by an analytical least squares fit
+			Z * A = F
+			Z.T * Z * A = Z.T * F
+			A = (Z.T * Z)^-1 * Z.T * F
+
+			**ARGUMENTS:** 
+
+			===== ==========================================     
+			data  array; the data that needs fitting
+			x     array; the x-coordinates of the data
+			y 	  array; the y-coordinates of the data
+			modes integer; number of modes to be fitted     
+			===== ==========================================
+
+			**RETURNS**
+			===== ===========================================
+			_A    Array with values of the Zernike polynomial
+			===== ===========================================
+		'''
+		if x is None and y is None:
+			y,x = np.mgrid[-1:1:self.n*1j, -1:1:self.m*1j]
+		_F 		= data.flatten()[:,np.newaxis]
+		_x 		= x.flatten()
+		_y 		= y.flatten()
+		_x0 	= (np.max(_x) - np.min(_x)) / 2 + np.min(_x)
+		_y0 	= (np.max(_y) - np.min(_y)) / 2 + np.min(_y)
+		_rho   	= np.sqrt( ( _x - _x0 )**2 + ( _y - _y0 )**2 )
+		_theta 	= np.arctan2( ( _y - _y0 ), ( _x - _x0 ) )
+
+		_Zern = []	
+		for i in range(modes):
+			_n, _m 	= self.i_to_nm(i)
+			_Znm 	= self.zernike_polynomial_polar(_n,_m,_rho,_theta)
+			_Zern.append(_Znm)
+			_Znm 	= None
+		_Zern 			= np.asarray(_Zern).T
+		_Zern_transpose = _Zern.T
+		_pseudo_inv 	= np.linalg.pinv(np.dot(_Zern_transpose, _Zern))
+		_int1 			= np.dot(_pseudo_inv, _Zern_transpose)
+		_A 				= np.dot(_int1,_F)
+		self.zernikeCoefficients = _A[:,0]
+		return _A[:,0]
+
+
+	def scale_zernike_coefficients(self, 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))
+		self.zernikeCoefficients 	= _c
+		return _c
+
+
+	def power_spectral_density(self, s0, p, rho_c, dx, dy, A):
+		# Lorentzian equation used as PSD function
+		m = self.m
+		n = self.n
+		UmList = [(M - ( m + 1) / 2) / (m * dx) for M in range(m)]
+		VnList = [(N - ( n + 1) / 2) / (n * dy) for N in range(n)]
+		VnList, UmList = np.meshgrid(UmList, VnList)
+
+		_rho_mn_list = np.sqrt(UmList**2 + VnList**2)
+		_h0list = 1 / ( 1 + np.power( ( _rho_mn_list / rho_c ), p) )
+		_h0_norm = np.sum(_h0list)
+		_rho_mn_list = np.asarray(_rho_mn_list).reshape(m,n)
+		
+		_PSD = (s0**2 * A) / (_h0_norm * (1 + np.power((_rho_mn_list/rho_c), p)))
+		self.PSD2D = _PSD
+		
+		return _PSD, _rho_mn_list
+
+
+	def phase_map(self, phase):
+		# Uniform distribution of the phase
+		_phase_in	 		= pyfftw.empty_aligned((self.n,self.m), dtype='complex128')
+		_phase_in[:] 		= phase
+		_phase_out	 		= pyfftw.empty_aligned((self.n, self.m), dtype='complex128')
+		_FFTW_Obj 	 		= pyfftw.FFTW( _phase_in, _phase_out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
+		_FFTW_Obj()	
+		_phaseAngle 		= np.angle(_phase_out)
+
+		return _phaseAngle
+
+
+	def generate_height_map(self, s0, p, rho_c, addToData=False):
+		'''
+		This function generates a surface height distribution based
+		on a Lorentzian PSD profile. Following the paper by E. Sidick (2009).
+
+		2020-01-20: Removed "mean=1e-4" from function def. --> Did this do anything?
+		'''
+
+		# Calculation of the necessary constants
+		_dx 				= 1 / self.m
+		_dy 				= 1 / self.n
+		_A 					= (self.m * _dx) * (self.n * _dy)
+		_phase 				= np.random.uniform(0,1,(self.m,self.m))
+		
+		# Calculating the PSD and phase
+		_PSD, _rho_mn_list 	= self.power_spectral_density(s0, p, rho_c, _dx, _dy, _A)
+		_phiMap 			= self.phase_map(_phase)
+
+		# Shifting the coordinates for the FFT
+		_PSD 				= np.fft.ifftshift(_PSD)
+
+		# Doing the fourier transform
+		_in	 				= pyfftw.empty_aligned((self.n,self.m), dtype='complex128')
+		_in[:] 				= np.exp( 1j * _phiMap) * np.sqrt(_A * np.abs(_PSD))
+		_out	 			= pyfftw.empty_aligned((self.n,self.m), dtype='complex128')
+		_FFTW_Obj 	 		= pyfftw.FFTW( _in, _out, axes=(0,1), direction='FFTW_BACKWARD', threads=6, flags=('FFTW_ESTIMATE',))
+		_FFTW_Obj()	
+
+		# The final height distribution
+		_hMap 				= 1 / ( _dx * _dy ) * _out
+		if not addToData:
+			self.data 		= _hMap.real
+		else:
+			self.data 		+= _hMap.real
+
+		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))
+		return self.data
+
+
+	def scale_rms(self, factor, data=None, filteredDataFlag=False):
+		# Scale the surface data such that the RMS gets multiplied by factor
+		_factor = factor ** 2
+		if data is None:
+			if filteredDataFlag:
+				self.scaledPSD2D    = self.filteredPSD2D * _factor
+				_data 				= self.filteredData
+			else:
+				self.calculate_2D_PSD()
+				self.scaledPSD2D    = self.PSD2D * _factor
+				_data 				= self.data
+		elif data is not None:
+			self.PSD2D 				= self.calculate_2D_PSD(data=data)
+			self.scaledPSD2D        = self.PSD2D * _factor
+			_data 					= data
+
+		# Fourier transform of the surface profile
+		_in	 				= pyfftw.empty_aligned((self.n,self.m), dtype='complex128')
+		_in[:] 				= _data
+		_out	 			= pyfftw.empty_aligned((self.n,self.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((self.n,self.m), dtype='complex128')
+		_in[:] 				= np.exp(1j*_phase) * np.sqrt(np.fft.fftshift(self.scaledPSD2D))
+		_out	 			= pyfftw.empty_aligned((self.n,self.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))
+		
+		return _out.real
+
+
+	def calculate_2D_PSD(self, data=None):
+		''' Returns the 2D PSD map '''
+		if data is None:
+			_in	 				= pyfftw.empty_aligned((self.n,self.m), dtype='complex128')
+			_in[:] 				= self.data
+			_out	 			= pyfftw.empty_aligned((self.n,self.m), dtype='complex128')
+		elif data is not None:
+			_n, _m 				= np.asarray(data).shape
+			_in	 				= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+			_in[:] 				= np.asarray(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()	
+		if data is None:
+			self.D1 = np.fft.fftshift(_out)
+			self.PSD2D = np.abs(self.D1)**2
+			return self.PSD2D
+		elif data is not None:
+			_D1 = np.fft.fftshift(_out)
+			_PSD2D = np.abs(_D1)**2
+			return _PSD2D
+
+
+	def calculate_1D_PSD(self, data=None, filteredDataFlag=False, rmax=None):
+		''' Returns the azimuthally averaged PSD map '''
+		if data is None:
+			try:
+				if not filteredDataFlag:
+					self.PSD1D = self.azimuthal_average(rmax)
+				else:
+					self.filteredPSD1D = self.azimuthal_average(filteredDataFlag=filteredDataFlag)
+					return self.filteredPSD1D
+			except AttributeError:
+				self.calculate_2D_PSD(data)
+				self.PSD1D = self.azimuthal_average(rmax)
+			return self.PSD1D
+		elif data is not None:
+			_PSD2D = self.calculate_2D_PSD(data=data)
+			_PSD1D = self.azimuthal_average(PSD=_PSD2D, rmax=rmax)
+			return _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
+		_x, _y 					= np.mgrid[-self.n/2*dpix : self.n/2*dpix : self.n*1j, -self.m/2*dpix : self.m/2*dpix : self.m*1j]
+		_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
+		return self.filteredData
+		
+
+	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]
+		_x, _y 					= np.mgrid[-self.n/2*dpix : self.n/2*dpix : self.n*1j, -self.m/2*dpix : self.m/2*dpix : self.m*1j]
+		_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
+		return self.filteredData
+		
+
+	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]
+		_x, _y 					= np.mgrid[-self.n/2*dpix : self.n/2*dpix : self.n*1j, -self.m/2*dpix : self.m/2*dpix : self.m*1j]
+		_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
+		return self.filteredData
+		
+
+	def azimuthal_average(self, PSD=None, center=None, filteredDataFlag=False, rmax=None):
+		"""
+		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 and PSD is None:
+			_PSD = self.PSD2D
+
+		elif PSD is None:
+			_PSD = self.filteredPSD2D
+
+		elif PSD is not None:
+			_PSD = PSD
+
+		if rmax == True:
+			_PSD[ _r > _x.max()/2 ] = np.nan
+		elif rmax:
+			_PSD[ _r > rmax ] = np.nan
+		
+		_i_sorted = _PSD.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
+
+	def generate_sine_pattern(self, ampx, ampy, perx, pery):
+		_x, _y = np.mgrid[-1:1:self.n*1j, -1:1:self.m*1j]
+		_z = ampx * np.sin(2*np.pi*_x*perx/2) + ampy * np.sin(2 * np.pi * _y * pery/2)
+		return _z
+
+	def generate_tilt(self, x, y, tx, ty):	
+		_z = y * np.tan(np.deg2rad(tx)) + x * np.tan(np.deg2rad(ty))
+		return _z

test cases (standalone)/fourierOptics_01_diffracting_aperture.py

@@ -0,0 +1,735 @@
+#!/usr/bin/env Python3
+import numpy as np
+from pdb import set_trace
+import matplotlib.pyplot as plt
+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 time
+
+
+class FourierOptics:
+	def __init__(self, l):
+		self.l = l # Wavelength
+
+
+	def planeWavePropagationFFTW(self, A, dpix=None, dz=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.abs(_E_xy0)**2).sum()	# Factor to scale the intensity after the FFTs
+		_IntensityNormFactor = np.sum(np.abs(_E_xy0)**2)
+		_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)
+		_kz = self.phaseFactorExponent(_fx, _fy, self.l) 	# origin at start of array
+		_phaseFactor = np.exp(-1j * _kz * dz) 				# origin at start of array
+		
+		# Calculating the spatial frequencies in k-space
+		_kx 		= np.fft.fftshift( _fx * 2 * np.pi )
+		_ky 		= np.fft.fftshift( _fy * 2 * np.pi )
+
+		_E_xyz 		= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_A_in 		= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_A_in[:] 	= _Az0 * _phaseFactor
+
+		_FFTW_Obj2 	= pyfftw.FFTW( _A_in, _E_xyz, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
+		_FFTW_Obj2()
+		_E_xyz 		= np.fft.fftshift(_E_xyz) 				# use fftshift to move the origin from edge to center
+		_E_xyz /= np.sqrt(_IntensityNormFactor) 			# Normalization back to initial energy.
+
+		# Coordinates of E_xy0
+		_y, _x = np.mgrid[-_n / 2 * dpix : _n / 2 * dpix : 1j*_n, -_m / 2 * dpix : _m / 2 * dpix : 1j * _m]
+		if outputLimits:
+			# 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.
+			_x				= np.linspace(-_n/2*dpix, _n/2*dpix, _n)
+			_y				= np.linspace(-_m/2*dpix, _m/2*dpix, _m)
+			_f_ip_real		= interpolate.interp2d(_x, _y, _E_xyz.real, kind='linear')
+			_f_ip_imag 		= interpolate.interp2d(_x, _y, _E_xyz.imag, kind='linear')
+			_xi				= np.linspace(outputLimits[0],outputLimits[1],_n)
+			_yi 			= np.linspace(outputLimits[0],outputLimits[1],_m)
+			_E_out_ip_real 	= _f_ip_real(_xi, _yi)
+			_E_out_ip_imag	= _f_ip_imag(_xi, _yi)
+			_E_xyz 			= _E_out_ip_real + _E_out_ip_imag*1j
+			_y, _x 			= np.mgrid[outputLimits[0]:outputLimits[1]:1j*_n, outputLimits[0]:outputLimits[1]:1j * _m]
+			
+
+		_E_xyz_extent = [np.min(_x), np.max(_x), np.min(_y), np.max(_x)]
+
+
+
+		if show or save:
+			fig = plt.figure(figsize=(8,3))
+			ax = fig.add_subplot(121)
+			im = ax.imshow(np.abs(_E_xyz), cmap='hot', extent=_E_xyz_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_xyz), cmap='hot', extent=_E_xyz_extent)
+			ax2.set_xlabel('x [mm]')
+			ax2.set_ylabel('y [mm]')
+			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_xyz, _x, _y])
+
+
+	def FresnelPropagationFFTW(self, A, xin, yin, dz=None, oversample=None, show=False, save=False, scaleInputArray=False, outputCoordinates='auto'):
+		'''
+			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]
+											If 'auto': the output coordinates will be determined 
+											automatically.
+
+			Output:
+			E_out 				- 	[array] (Complex) Output field distribution
+			x, y 				-	[array] Output coordinates 
+		'''
+		assert np.shape(xin) == np.shape(A), 'Input coordinates (xin) do not match shape of input field.'
+		assert np.shape(yin) == np.shape(A), 'Input coordinates (yin) do not match shape of input field.'
+		A = np.conj(A)
+
+		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='linear')
+			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)
+		_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
+		
+		# Calculate the pixel size from the xi, eta plane
+		_dpix = np.mean(np.diff(_xi))
+		if _dpix < 1e-99:
+			_dpix = np.mean(np.diff(_xi.T))
+
+		# Gaussian test: use 8*w0 radius for FFT:
+		if outputCoordinates == 'auto':
+			_w0_fit, _ 		= leastsq(self.minimize, self.l, args=(0, _xi, _eta, A))
+			_gauss 			= GaussianBeam(_w0_fit, self.l)
+			_halfwidth 		= 8 * _gauss.beamRadius(dz)[0]
+			_xip			= np.linspace(-_halfwidth, _halfwidth, _n) 						# Desired x,y - plane
+			_yip			= np.linspace(-_halfwidth, _halfwidth, _m) 						# Desired x,y - plane
+		else:
+			_xip			= np.linspace(outputCoordinates[0], outputCoordinates[1], _n) 	# Desired x,y - plane
+			_yip			= np.linspace(outputCoordinates[2], outputCoordinates[3], _m) 	# Desired x,y - plane
+
+
+		# The quadratic phase exponential of the input wavefront
+		_phaseComponent = np.exp(1j * _k / (2 * dz) * (_xi**2 + _eta**2))
+
+		# Complete input
+		_U_in 			= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_U_in[:]		= A * _phaseComponent
+		
+		# Output of the FFT evaluated at frequency fx = x / (lambda * z)
+		_U_out 			= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_FFTW_Obj		= pyfftw.FFTW( np.fft.ifftshift(_U_in), _U_out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
+		_FFTW_Obj()
+		_U_out 			= np.fft.fftshift(_U_out)
+
+		# Evaluate frequencies:
+		_fx 			= np.fft.fftshift(np.fft.fftfreq(np.shape(_xi)[0], d=_dpix))
+		_fy 			= np.fft.fftshift(np.fft.fftfreq(np.shape(_xi)[1], d=_dpix))
+		_x 				= _fx * self.l * dz
+		_y 				= _fy * self.l * dz
+		_x2, _y2 		= np.meshgrid(_x, _y)
+
+
+		# Final output and frequencies
+		_pre_amble 		= np.exp(1j * _k * dz) / (1j * self.l * dz) * np.exp(1j * _k / (2 * dz) * (_x2**2 + _y2**2))
+		_E_out 			= _pre_amble * _U_out
+
+		# Normalization of the output values
+		_IntensityNormFactor = np.sum(np.abs(A)**2)
+		_E_out /= (np.abs(np.exp(1j*_k*dz)/(1j*self.l*dz)) * np.sqrt(_IntensityNormFactor) * np.sqrt(_n*_m))
+
+		# Performing interpolation to get to the required final plane, done for real and imaginary part.
+		_f_ip_real		= interpolate.interp2d(_x, _y, _E_out.real, kind='linear')
+		_f_ip_imag 		= interpolate.interp2d(_x, _y, _E_out.imag, kind='linear')
+		_E_out_ip_real 	= _f_ip_real(_xip, _yip)
+		_E_out_ip_imag	= _f_ip_imag(_xip, _yip)
+		_E_out_ip 		= _E_out_ip_real + _E_out_ip_imag*1j
+		_E_out_ip 		/= np.sqrt(((_x2.max()/_xip.max())*(_y2.max()/_yip.max())))
+		_E_out_ip = np.conj(_E_out_ip)
+
+		if show:
+			# Plotting the results
+			_fig = plt.figure(figsize=(8,3))
+			_ax = _fig.add_subplot(121)
+			_im = _ax.imshow(np.abs(_E_out_ip), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()])
+			_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='Amplitude')
+			_ax2 = _fig.add_subplot(122)
+			_im2 = _ax2.imshow(np.angle(_E_out_ip), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()])
+			_ax2.set_xlim([_xip.min(), _xip.max()])
+			_ax2.set_ylim([_yip.min(), _yip.max()])
+			_ax2.set_xlabel('x [mm]')
+			_ax2.set_ylabel('y [mm]')
+			_fig.colorbar(_im2, ax=_ax2, label='Phase')
+			plt.tight_layout()
+			if save:
+				plt.savefig(save, bbox_inches='tight', transparent=True)
+			plt.show()
+
+		_x, _y = np.meshgrid(_xip, _yip)
+		return np.array([_E_out_ip, _x, _y])
+
+
+	def freqArray(self, A, dpix):
+		n, m = np.shape(A)
+		
+		fy, fx = np.mgrid[0:m, 0:n]
+
+		fx = fx.flatten()
+		fy = fy.flatten()
+		
+		for i in range(len(fx)):
+			fx[i] = (fx[i] - n*bool(fx[i]>(n-1)/2) ) / (n*dpix)
+			fy[i] = (fy[i] - m*bool(fy[i]>(m-1)/2) ) / (m*dpix)
+		
+		fx = fx.reshape(n,m)
+		fy = fy.reshape(n,m)
+		
+		return fx, fy
+				
+
+	def phaseFactorExponent(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 minimize(self, a, dz, x, y, og_data):
+		'''
+			Residual function used for fitting the gaussian beam to the input distribution.
+			Input:
+			a 		- 	[float] Guess for w0
+			dz  	- 	[float] Propagation distance, should be known
+			x,y 	- 	[array] meshgrid of the x and y coordinates.
+			og_data - 	[array] input distribution to be fitted.
+		'''
+		return np.abs((self.gaussianBeamFit(a,dz,x,y) - og_data).flatten()).real
+
+
+	def gaussianBeamFit(self, a, dz, x, y):
+		'''
+			Function used to fit the Gaussian beam to an input plane. Very similar
+			to the GaussianBeam class below, but modified to allow a best fit.
+			Input:
+			a 	- 	[float] Guess for w0
+			dz  - 	[float] Propagation distance, should be known
+			x,y - 	[array] meshgrid of the x and y coordinates.
+		'''
+		_w0 = a
+		_z 	= dz
+		_r 	  	= np.sqrt(x**2 + y**2)
+		if _z != 0:
+			_z_c 	= np.pi * _w0**2 / self.l
+			_w 		= _w0 * np.sqrt(1 + (_z/_z_c)**2)
+			_R 		= _z + _z_c **2 / _z
+			_phi0 	= np.arctan(_z/_z_c)
+			_k 		= 2 * np.pi / self.l
+			_results = np.sqrt(2 / (np.pi*_w**2)) * np.exp(- _r**2 / _w**2 - 1j * _k * _z - 1j * np.pi * _r**2 / (self.l * _R) + 1j*_phi0)
+		elif _z == 0:
+			_results = np.sqrt(2 / (np.pi*_w0**2)) * np.exp( - _r**2 / _w0**2 )
+		return _results
+
+
+
+
+
+
+class profiles:
+	def __init__(self):
+		pass
+
+	def topHat(self, r, n, m, dpix, oversample=1):
+	    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 <= (r)).astype(int)
+	    return _Anew
+
+	def obscuredTopHat(self, r1, r2, n, m, dpix, oversample=1):
+		'''
+			Creates a grid of n x m points, following a top hat function
+			where all points inside a radius r is 1 and 0 otherwise.
+		'''
+		assert r1 > r2, "r2 < r1, change input"
+		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))
+		return _Anew.astype(int)
+
+
+	def spiderTopHat(self, r1, r2, n, m, dpix, oversample=1):
+		'''
+			Creates a grid of n x m points, following a top hat function
+			where all points inside a radius r is 1 and 0 otherwise.
+		'''
+		assert r1 > r2, "r2 < r1, change input"
+		j,i 				= np.mgrid[-n/2*dpix : n/2*dpix : 1j*n*oversample, -m/2*dpix : m/2*dpix : 1j*m*oversample]
+		n *= oversample
+		m *= oversample
+		dist_from_center 	= np.sqrt(i**2 + j**2)
+		_Anew 				= (dist_from_center <= (r1))  & (dist_from_center >= (r2))
+		_Aspider 			= (np.abs(i) < (0.05*r1)) | (np.abs(j) < (0.05*r1))
+		_Anew 				= _Anew & ~_Aspider
+		return _Anew
+
+
+
+class GaussianBeam:
+	def __init__(self, w0, l):
+		# Confocal distance
+		self.z_c 	 = np.pi * w0**2 / l
+
+		# Far-field divergence angle
+		self.theta_0 = l / (np.pi * w0)
+		self.l 		 = l
+		self.w0 	 = w0
+		self.fwhm 	 = 1.18 * self.theta_0
+
+	def fundamentalMode(self, z, n, m, rmax):
+		y, x = np.mgrid[-rmax:rmax:n*1j,-rmax:rmax:n*1j]
+		r 	 = np.sqrt(x**2 + y**2)
+		if z != 0:
+			_w 		= self.beamRadius(z)
+			_R 		= self.curvature(z)
+			_phi0 	= self.phaseShift(z)
+			_k 		= 2 * np.pi / self.l
+			return np.sqrt(2 / (np.pi*_w**2)) * np.exp(- r**2 / _w**2 - 1j * _k * z - 1j * np.pi * r**2 / (self.l * _R) + 1j*_phi0)
+		elif z == 0:
+			return np.sqrt(2 / (np.pi*self.w0**2)) * np.exp( - r**2 / self.w0**2 ) 
+
+	def curvature(self, z):
+		return z + self.z_c **2 / z
+
+	def beamRadius(self, z):
+		return self.w0 * np.sqrt(1 + (z/self.z_c)**2)
+
+	def phaseShift(self, z):
+		return np.arctan(z/self.z_c)
+
+	def relativePower(self, r, z):
+		return np.exp(-2 * (r / self.beamRadius(z))**2)
+
+
+
+class lensObject:
+	'''
+		The lens object describes any spherical lens and find the phase delay
+		for all light traveling through it.
+		The refractive index, of course, depends on the wavelength.
+	'''
+	def __init__(self, glass, R1, R2, thickness, CA, f=None, name=None):
+		self.glass 		= glass
+		assert glassCatalog.catalog.get(self.glass), 'Glass type not in catalog'
+		self.R1 		= R1 				# mm
+		self.R2 		= R2 				# mm
+		self.thickness 	= thickness			# mm
+		self.CA 		= CA 				# mm
+		self.dispersionCoefficients = glassCatalog.catalog.get(self.glass)
+		self.referenceWavelength = 1.35 	# micron
+		self.type 		= 'Spherical lens'
+		if not f:
+			self.f 		= self.findFocalLength()
+		else:
+			self.f 		= f
+		if name:
+			assert type(name)==str, 'No string for a name.'
+			self.name = name
+		else:
+			self.name = str(self)
+
+	def findRefractiveIndex(self, _l):
+		_a = self.dispersionCoefficients
+		_n = np.sqrt( (_a[0]* _l**2)/(_l**2 - _a[1]) + (_a[2] * _l**2) / (_l**2 - _a[3]) + (_a[4] * _l**2) / (_l**2 - _a[5]) + 1)
+		return _n
+
+	def findFocalLength(self):
+		_n 		= self.findRefractiveIndex(self.referenceWavelength)
+		_f_recip = (_n - 1) * (1 / self.R1 - 1 / self.R2)
+		return 1/_f_recip
+
+	def phaseDelayFunction(self, x, y, l, paraxial=False):
+		if not paraxial:
+			_A1 = self.R1 * (1 - np.sqrt(1 - (x**2 + y**2) / self.R1**2 ) )
+			_A2 = self.R2 * (1 - np.sqrt(1 - (x**2 + y**2) / self.R2**2 ) )
+			_dxy = self.thickness - _A1 + _A2
+		if paraxial:
+			_dxy = self.thickness - (x**2 + y**2)/2 * (1/self.R1 - 1/self.R2)
+		_k = 2 * np.pi / 30e-3
+		_n = self.findRefractiveIndex(self.referenceWavelength)
+		_t_l = np.exp(-1j * _k * self.thickness)* np.exp(-1j * _k * (_n - 1) * _dxy)
+
+		# 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)).astype(int)
+		_t_l *= _lensMask
+		return _t_l
+
+
+
+class mirrorObject:
+	'''
+		Describes any mirror. No wavelength dependence exist. 
+		Coatings are not considered.
+	'''
+	def __init__(self, dx, dy, name=None):
+		self.dx = dx
+		self.dy = dy
+		if name:
+			assert type(name)==str, 'No string for a name.'
+			self.name = name
+		else:
+			self.name = str(self)
+
+
+class parabolicMirror(mirrorObject):
+	def __init__(self, f, CA, dx=None, dy=None, name=None):
+		mirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.f = f
+		self.CA = CA
+		self.type = 'Parabolic mirror'
+
+	def phaseDelayFunction(self, x, y, l):
+		_k 		= 2 * np.pi / l
+		_dz 	= (x**2 + y**2) / (4 * self.f)
+		_t_m 	= np.exp(-1j * _k * 2 * _dz)
+		return _t_m
+
+
+
+class sphericalMirror(mirrorObject):
+	def __init__(self, R, CA, dx=None, dy=None, name=None):
+		mirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.R = R
+		self.CA = CA
+		self.type = 'Spherical mirror'
+
+	def phaseDelayFunction(self, x, y, l):
+		_k = 2 * np.pi / l
+		_dz = self.R * (1 - np.sqrt(1 - (x**2 + y**2)/self.R**2) )
+		_t_m = np.exp(-1j * _k * 2 * _dz)
+		return _t_m
+
+
+
+class freeformMirror(mirrorObject):
+	def __init__(self, x, y, z, dx=None, dy=None, name=None):
+		mirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.x = x
+		self.y = y
+		self.z = z
+		self.type = 'User def. mirror'
+
+	def phaseDelayFunction(self, l):
+		_k = 2 * np.pi / l
+		_t_m = np.exp(-1j * _k * 2 * self.z)
+		return _t_m
+
+
+
+
+class OpticalTrain:
+	'''
+		The main class where an optical system is build up sequentially.
+	'''
+	def __init__(self):
+		self.opticalTrain 		= []
+		self.sequentialDistance = []
+
+	def add_object(self, obj, dist):
+		'''
+			Adding a lens or mirror object to the opticalTrain
+			obj 	The lens or mirror object
+			dist 	The distance from the the previous object.
+		'''
+		assert type(dist) == float or type(dist) == int, 'Distance is not a float or integer.'
+		self.opticalTrain.append(obj)
+		self.sequentialDistance.append(float(dist))
+
+	def print_lens_Data(self):
+		'''
+			Print the data of the current opticalTrain.
+		'''
+		_tableData = []
+		for _lens in self.opticalTrain:
+			_i = self.opticalTrain.index(_lens)
+			_distances = self.sequentialDistance
+			_distances.append('-')
+			_tableData.append([_i, _lens.name, _lens.type, _distances[_i+1] ])
+		print(tabulate(_tableData, headers=['#', 'Name', 'Type', 'Distance to next surface']))
+
+
+	def propagate_system(self, U_in, x, y, l, output=None, paraxial=True):
+		# We need to have an input distribution
+		# We need to determine what propagation method is required
+		# Call the objects in the self.opticalTrain sequentially
+		# Call the subsequent phaseDelay
+		if output == 'complete':
+			_wf_list 	= []
+			_x_list 	= []
+			_y_list 	= []
+			_wf_list.append(_wf)
+			_x_list.append(x)
+			_y_list.append(y)
+		
+		_wf = U_in
+
+		for _lens in self.opticalTrain:
+			_i 	= self.opticalTrain.index(_lens)
+			_dz = self.sequentialDistance[_i]
+			_wf = _lens.phaseDelayFunction(x, y, l, paraxial=paraxial)
+
+			# Perform Fresnel number test:
+			_F = self.fresnel_number()
+
+			if _F < 5:
+				# Fresnel propagation
+				_wf, x, y = self.fresnel_propagation()
+			elif _F >= 5:
+				# Angular spectrum (p2p) propagation
+				_wf, x, y = self.angular_spectrum_propagation()
+
+			if output == 'complete':
+				_wf_list.append(_wf)
+				_x_list.append(x)
+				_y_list.append(y)
+		
+		if output == 'Complete':
+			return np.array([_wf_list, _x_list, _y_list])
+
+		else:
+			return np.array([_wf, x, y])
+
+		
+
+
+
+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*l
+	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.spiderTopHat(250, 100, 2001, 2001, dpix=1, oversample=1)
+	A = prof.topHat(4*l, n, m, dpix=dpix, oversample=1)
+
+	# Alternatively, create a gaussian profile
+	# gauss = GaussianBeam(0.3, l*1000)
+	# A = gauss.fundamentalMode(0,n,m,np.max(x))
+
+	# 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))
+		
+	# The input distribution gets a curved wavefront after the lens
+	U_in = A
+	
+
+	Fresnel = False
+	P2P = not Fresnel
+
+	oC = [x.min(), x.max(), y.min(), y.max()]
+	# Performing the propagation using Fresnel approximation to the focal plane.
+	if Fresnel:
+		savename = 'fourierOptics_01_diffracting_aperture_fresnel.pdf'
+		Exyz,_,_ = fou.FresnelPropagationFFTW(U_in, x, y, dz=4*l, oversample=None, show=True, save=savename, outputCoordinates=oC)
+	if P2P:
+		savename = 'fourierOptics_01_diffracting_aperture_p2p.pdf'
+		Exyz,_,_ = fou.planeWavePropagationFFTW(A, dz=4*l, dpix=dpix, oversample=None, show=True, save=savename)
+	
+	# 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(x), np.max(x),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') 
+	if Fresnel:
+		plt.savefig('fourierOptics_01_diffracting_aperture_fresnel_slice.pdf', bbox_inches='tight', transparent=True)
+	if P2P:
+		plt.savefig('fourierOptics_01_diffracting_aperture_p2p_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 True:
+		Ics 	= []
+		Ics_x 	= []
+		Isum 	= []
+		xmax = []
+		zList = np.linspace(1,10.1,18).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
+			if Fresnel:
+				Exyz, x, y = fou.planeWavePropagationFFTW(A, dz=i*l, dpix=dpix, oversample=None, show=False, save=False)#, scaleInputArray=(n,m))
+			if P2P:
+				Exyz, x, y = fou.FresnelPropagationFFTW(U_in, xin, yin, dz=i*l, oversample=None, 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(211)
+		ax.tick_params(axis ='both', direction='in') 
+
+		# Plotting the square root of the amplitude for better contrast
+		im = ax.imshow(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'))
+		ax.set_ylim([-0.2, 0.2])
+		
+		# Plotting the intensity relative to the input
+		ax2 = fig.add_subplot(212)
+		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 ($\\times \lambda$)')
+		ax2.tick_params(axis ='both', direction='in') 
+		ax2.set_ylabel('Intensity')
+		ax2.set_xlim([np.min(zList), np.max(zList)])
+		ax2.grid(True, alpha=0.3)
+
+		# Saving the figure? Then show.
+		plt.tight_layout()
+		# plt.savefig('fresnel_cross.pdf', bbox_inches='tight')
+		# plt.savefig('fresnel_cross.png', bbox_inches='tight')
+		if P2P:
+			plt.savefig('fourierOptics_01_diffracting_aperture_p2p_cross.pdf', bbox_inches='tight', transparent=True)
+		if Fresnel:
+			plt.savefig('fourierOptics_01_diffracting_aperture_fresnel_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)))
+	
+
+

test cases (standalone)/fourierOptics_025_single_lens_system.py

@@ -0,0 +1,729 @@
+#!/usr/bin/env Python3
+import numpy as np
+from pdb import set_trace
+import matplotlib.pyplot as plt
+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 time
+
+
+class FourierOptics:
+	def __init__(self, l):
+		self.l = l # Wavelength
+
+
+	def planeWavePropagationFFTW(self, A, dpix=None, dz=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.abs(_E_xy0)**2).sum()	# Factor to scale the intensity after the FFTs
+		_IntensityNormFactor = np.sum(np.abs(_E_xy0)**2)
+		_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)
+		_kz = self.phaseFactorExponent(_fx, _fy, self.l) 	# origin at start of array
+		_phaseFactor = np.exp(-1j * _kz * dz) 				# origin at start of array
+		
+		# Calculating the spatial frequencies in k-space
+		_kx 		= np.fft.fftshift( _fx * 2 * np.pi )
+		_ky 		= np.fft.fftshift( _fy * 2 * np.pi )
+
+		_E_xyz 		= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_A_in 		= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_A_in[:] 	= _Az0 * _phaseFactor
+
+		_FFTW_Obj2 	= pyfftw.FFTW( _A_in, _E_xyz, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
+		_FFTW_Obj2()
+		_E_xyz 		= np.fft.fftshift(_E_xyz) 				# use fftshift to move the origin from edge to center
+		_E_xyz /= np.sqrt(_IntensityNormFactor) 			# Normalization back to initial energy.
+
+		# Coordinates of E_xy0
+		_y, _x = np.mgrid[-_n / 2 * dpix : _n / 2 * dpix : 1j*_n, -_m / 2 * dpix : _m / 2 * dpix : 1j * _m]
+		if outputLimits:
+			# 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.
+			_x				= np.linspace(-_n/2*dpix, _n/2*dpix, _n)
+			_y				= np.linspace(-_m/2*dpix, _m/2*dpix, _m)
+			_f_ip_real		= interpolate.interp2d(_x, _y, _E_xyz.real, kind='linear')
+			_f_ip_imag 		= interpolate.interp2d(_x, _y, _E_xyz.imag, kind='linear')
+			_xi				= np.linspace(outputLimits[0],outputLimits[1],_n)
+			_yi 			= np.linspace(outputLimits[0],outputLimits[1],_m)
+			_E_out_ip_real 	= _f_ip_real(_xi, _yi)
+			_E_out_ip_imag	= _f_ip_imag(_xi, _yi)
+			_E_xyz 			= _E_out_ip_real + _E_out_ip_imag*1j
+			_y, _x 			= np.mgrid[outputLimits[0]:outputLimits[1]:1j*_n, outputLimits[0]:outputLimits[1]:1j * _m]
+			
+
+		_E_xyz_extent = [np.min(_x), np.max(_x), np.min(_y), np.max(_x)]
+
+
+
+		if show or save:
+			fig = plt.figure(figsize=(8,3))
+			ax = fig.add_subplot(121)
+			im = ax.imshow(np.abs(_E_xyz), cmap='hot', extent=_E_xyz_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_xyz), cmap='hot', extent=_E_xyz_extent)
+			ax2.set_xlabel('x [mm]')
+			ax2.set_ylabel('y [mm]')
+			fig.colorbar(im2, ax=ax2)
+			plt.tight_layout()
+		if save:
+			plt.savefig(save, bbox_inches='tight')
+		if show:
+			plt.show()
+		elif not show:
+			plt.close()
+
+		return np.array([_E_xyz, _x, _y])
+
+
+	def FresnelPropagationFFTW(self, A, xin, yin, dz=None, oversample=None, show=False, save=False, scaleInputArray=False, outputCoordinates='auto'):
+		'''
+			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]
+											If 'auto': the output coordinates will be determined 
+											automatically.
+
+			Output:
+			E_out 				- 	[array] (Complex) Output field distribution
+			x, y 				-	[array] Output coordinates 
+		'''
+		assert np.shape(xin) == np.shape(A), 'Input coordinates (xin) do not match shape of input field.'
+		assert np.shape(yin) == np.shape(A), 'Input coordinates (yin) do not match shape of input field.'
+		A = np.conj(A)
+
+		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='linear')
+			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)
+		_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
+		
+		# Calculate the pixel size from the xi, eta plane
+		_dpix = np.mean(np.diff(_xi))
+		if _dpix < 1e-99:
+			_dpix = np.mean(np.diff(_xi.T))
+
+		# Gaussian test: use 8*w0 radius for FFT:
+		if outputCoordinates == 'auto':
+			_w0_fit, _ 		= leastsq(self.minimize, self.l, args=(0, _xi, _eta, A))
+			_gauss 			= GaussianBeam(_w0_fit, self.l)
+			_halfwidth 		= 8 * _gauss.beamRadius(dz)[0]
+			_xip			= np.linspace(-_halfwidth, _halfwidth, _n) 						# Desired x,y - plane
+			_yip			= np.linspace(-_halfwidth, _halfwidth, _m) 						# Desired x,y - plane
+		else:
+			_xip			= np.linspace(outputCoordinates[0], outputCoordinates[1], _n) 	# Desired x,y - plane
+			_yip			= np.linspace(outputCoordinates[2], outputCoordinates[3], _m) 	# Desired x,y - plane
+
+
+		# The quadratic phase exponential of the input wavefront
+		_phaseComponent = np.exp(1j * _k / (2 * dz) * (_xi**2 + _eta**2))
+
+		# Complete input
+		_U_in 			= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_U_in[:]		= A * _phaseComponent
+		
+		# Output of the FFT evaluated at frequency fx = x / (lambda * z)
+		_U_out 			= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_FFTW_Obj		= pyfftw.FFTW( np.fft.ifftshift(_U_in), _U_out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
+		_FFTW_Obj()
+		_U_out 			= np.fft.fftshift(_U_out)
+
+		# Evaluate frequencies:
+		_fx 			= np.fft.fftshift(np.fft.fftfreq(np.shape(_xi)[0], d=_dpix))
+		_fy 			= np.fft.fftshift(np.fft.fftfreq(np.shape(_xi)[1], d=_dpix))
+		_x 				= _fx * self.l * dz
+		_y 				= _fy * self.l * dz
+		_x2, _y2 		= np.meshgrid(_x, _y)
+
+
+		# Final output and frequencies
+		_pre_amble 		= np.exp(1j * _k * dz) / (1j * self.l * dz) * np.exp(1j * _k / (2 * dz) * (_x2**2 + _y2**2))
+		_E_out 			= _pre_amble * _U_out
+
+		# Normalization of the output values
+		_IntensityNormFactor = np.sum(np.abs(A)**2)
+		_E_out /= (np.abs(np.exp(1j*_k*dz)/(1j*self.l*dz)) * np.sqrt(_IntensityNormFactor) * np.sqrt(_n*_m))
+
+		# Performing interpolation to get to the required final plane, done for real and imaginary part.
+		_f_ip_real		= interpolate.interp2d(_x, _y, _E_out.real, kind='linear')
+		_f_ip_imag 		= interpolate.interp2d(_x, _y, _E_out.imag, kind='linear')
+		_E_out_ip_real 	= _f_ip_real(_xip, _yip)
+		_E_out_ip_imag	= _f_ip_imag(_xip, _yip)
+		_E_out_ip 		= _E_out_ip_real + _E_out_ip_imag*1j
+		_E_out_ip 		/= np.sqrt(((_x2.max()/_xip.max())*(_y2.max()/_yip.max())))
+		_E_out_ip = np.conj(_E_out_ip)
+
+		if show:
+			# Plotting the results
+			_fig = plt.figure(figsize=(8,3))
+			_ax = _fig.add_subplot(121)
+			_im = _ax.imshow(np.abs(_E_out_ip), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()])
+			_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='Amplitude')
+			_ax2 = _fig.add_subplot(122)
+			_im2 = _ax2.imshow(np.angle(_E_out_ip), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()])
+			_ax2.set_xlim([_xip.min(), _xip.max()])
+			_ax2.set_ylim([_yip.min(), _yip.max()])
+			_ax2.set_xlabel('x [mm]')
+			_ax2.set_ylabel('y [mm]')
+			_fig.colorbar(_im2, ax=_ax2, label='Phase')
+			plt.tight_layout()
+			if save:
+				plt.savefig(save, bbox_inches='tight', transparent=True)
+			plt.show()
+
+		_x, _y = np.meshgrid(_xip, _yip)
+		return np.array([_E_out_ip, _x, _y])
+
+
+	def freqArray(self, A, dpix):
+		n, m = np.shape(A)
+		
+		fy, fx = np.mgrid[0:m, 0:n]
+
+		fx = fx.flatten()
+		fy = fy.flatten()
+		
+		for i in range(len(fx)):
+			fx[i] = (fx[i] - n*bool(fx[i]>(n-1)/2) ) / (n*dpix)
+			fy[i] = (fy[i] - m*bool(fy[i]>(m-1)/2) ) / (m*dpix)
+		
+		fx = fx.reshape(n,m)
+		fy = fy.reshape(n,m)
+		
+		return fx, fy
+				
+
+	def phaseFactorExponent(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 minimize(self, a, dz, x, y, og_data):
+		'''
+			Residual function used for fitting the gaussian beam to the input distribution.
+			Input:
+			a 		- 	[float] Guess for w0
+			dz  	- 	[float] Propagation distance, should be known
+			x,y 	- 	[array] meshgrid of the x and y coordinates.
+			og_data - 	[array] input distribution to be fitted.
+		'''
+		return np.abs((self.gaussianBeamFit(a,dz,x,y) - og_data).flatten()).real
+
+
+	def gaussianBeamFit(self, a, dz, x, y):
+		'''
+			Function used to fit the Gaussian beam to an input plane. Very similar
+			to the GaussianBeam class below, but modified to allow a best fit.
+			Input:
+			a 	- 	[float] Guess for w0
+			dz  - 	[float] Propagation distance, should be known
+			x,y - 	[array] meshgrid of the x and y coordinates.
+		'''
+		_w0 = a
+		_z 	= dz
+		_r 	  	= np.sqrt(x**2 + y**2)
+		if _z != 0:
+			_z_c 	= np.pi * _w0**2 / self.l
+			_w 		= _w0 * np.sqrt(1 + (_z/_z_c)**2)
+			_R 		= _z + _z_c **2 / _z
+			_phi0 	= np.arctan(_z/_z_c)
+			_k 		= 2 * np.pi / self.l
+			_results = np.sqrt(2 / (np.pi*_w**2)) * np.exp(- _r**2 / _w**2 - 1j * _k * _z - 1j * np.pi * _r**2 / (self.l * _R) + 1j*_phi0)
+		elif _z == 0:
+			_results = np.sqrt(2 / (np.pi*_w0**2)) * np.exp( - _r**2 / _w0**2 )
+		return _results
+
+
+
+
+
+
+class profiles:
+	def __init__(self):
+		pass
+
+	def topHat(self, r, n, m, dpix, oversample=1):
+	    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 <= (r)).astype(int)
+	    return _Anew
+
+	def obscuredTopHat(self, r1, r2, n, m, dpix, oversample=1):
+		'''
+			Creates a grid of n x m points, following a top hat function
+			where all points inside a radius r is 1 and 0 otherwise.
+		'''
+		assert r1 > r2, "r2 < r1, change input"
+		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))
+		return _Anew.astype(int)
+
+
+	def spiderTopHat(self, r1, r2, n, m, dpix, oversample=1):
+		'''
+			Creates a grid of n x m points, following a top hat function
+			where all points inside a radius r is 1 and 0 otherwise.
+		'''
+		assert r1 > r2, "r2 < r1, change input"
+		j,i 				= np.mgrid[-n/2*dpix : n/2*dpix : 1j*n*oversample, -m/2*dpix : m/2*dpix : 1j*m*oversample]
+		n *= oversample
+		m *= oversample
+		dist_from_center 	= np.sqrt(i**2 + j**2)
+		_Anew 				= (dist_from_center <= (r1))  & (dist_from_center >= (r2))
+		_Aspider 			= (np.abs(i) < (0.05*r1)) | (np.abs(j) < (0.05*r1))
+		_Anew 				= _Anew & ~_Aspider
+		return _Anew
+
+
+
+class GaussianBeam:
+	def __init__(self, w0, l):
+		# Confocal distance
+		self.z_c 	 = np.pi * w0**2 / l
+
+		# Far-field divergence angle
+		self.theta_0 = l / (np.pi * w0)
+		self.l 		 = l
+		self.w0 	 = w0
+		self.fwhm 	 = 1.18 * self.theta_0
+
+	def fundamentalMode(self, z, n, m, rmax):
+		y, x = np.mgrid[-rmax:rmax:n*1j,-rmax:rmax:n*1j]
+		r 	 = np.sqrt(x**2 + y**2)
+		if z != 0:
+			_w 		= self.beamRadius(z)
+			_R 		= self.curvature(z)
+			_phi0 	= self.phaseShift(z)
+			_k 		= 2 * np.pi / self.l
+			return np.sqrt(2 / (np.pi*_w**2)) * np.exp(- r**2 / _w**2 - 1j * _k * z - 1j * np.pi * r**2 / (self.l * _R) + 1j*_phi0)
+		elif z == 0:
+			return np.sqrt(2 / (np.pi*self.w0**2)) * np.exp( - r**2 / self.w0**2 ) 
+
+	def curvature(self, z):
+		return z + self.z_c **2 / z
+
+	def beamRadius(self, z):
+		return self.w0 * np.sqrt(1 + (z/self.z_c)**2)
+
+	def phaseShift(self, z):
+		return np.arctan(z/self.z_c)
+
+	def relativePower(self, r, z):
+		return np.exp(-2 * (r / self.beamRadius(z))**2)
+
+
+
+class lensObject:
+	'''
+		The lens object describes any spherical lens and find the phase delay
+		for all light traveling through it.
+		The refractive index, of course, depends on the wavelength.
+	'''
+	def __init__(self, glass, R1, R2, thickness, CA, f=None, name=None):
+		self.glass 		= glass
+		assert glassCatalog.catalog.get(self.glass), 'Glass type not in catalog'
+		self.R1 		= R1 				# mm
+		self.R2 		= R2 				# mm
+		self.thickness 	= thickness			# mm
+		self.CA 		= CA 				# mm
+		self.dispersionCoefficients = glassCatalog.catalog.get(self.glass)
+		self.referenceWavelength = 1.35 	# micron
+		self.type 		= 'Spherical lens'
+		if not f:
+			self.f 		= self.findFocalLength()
+		else:
+			self.f 		= f
+		if name:
+			assert type(name)==str, 'No string for a name.'
+			self.name = name
+		else:
+			self.name = str(self)
+
+	def findRefractiveIndex(self, _l):
+		_a = self.dispersionCoefficients
+		_n = np.sqrt( (_a[0]* _l**2)/(_l**2 - _a[1]) + (_a[2] * _l**2) / (_l**2 - _a[3]) + (_a[4] * _l**2) / (_l**2 - _a[5]) + 1)
+		return _n
+
+	def findFocalLength(self):
+		_n 		= self.findRefractiveIndex(self.referenceWavelength)
+		_f_recip = (_n - 1) * (1 / self.R1 - 1 / self.R2)
+		return 1/_f_recip
+
+	def phaseDelayFunction(self, x, y, l, paraxial=False):
+		if not paraxial:
+			_A1 = self.R1 * (1 - np.sqrt(1 - (x**2 + y**2) / self.R1**2 ) )
+			_A2 = self.R2 * (1 - np.sqrt(1 - (x**2 + y**2) / self.R2**2 ) )
+			_dxy = self.thickness - _A1 + _A2
+		if paraxial:
+			_dxy = self.thickness - (x**2 + y**2)/2 * (1/self.R1 - 1/self.R2)
+		_k = 2 * np.pi / 30e-3
+		_n = self.findRefractiveIndex(self.referenceWavelength)
+		_t_l = np.exp(-1j * _k * self.thickness)* np.exp(-1j * _k * (_n - 1) * _dxy)
+
+		# 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)).astype(int)
+		_t_l *= _lensMask
+		return _t_l
+
+
+
+class mirrorObject:
+	'''
+		Describes any mirror. No wavelength dependence exist. 
+		Coatings are not considered.
+	'''
+	def __init__(self, dx, dy, name=None):
+		self.dx = dx
+		self.dy = dy
+		if name:
+			assert type(name)==str, 'No string for a name.'
+			self.name = name
+		else:
+			self.name = str(self)
+
+
+class parabolicMirror(mirrorObject):
+	def __init__(self, f, CA, dx=None, dy=None, name=None):
+		mirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.f = f
+		self.CA = CA
+		self.type = 'Parabolic mirror'
+
+	def phaseDelayFunction(self, x, y, l):
+		_k 		= 2 * np.pi / l
+		_dz 	= (x**2 + y**2) / (4 * self.f)
+		_t_m 	= np.exp(-1j * _k * 2 * _dz)
+		return _t_m
+
+
+
+class sphericalMirror(mirrorObject):
+	def __init__(self, R, CA, dx=None, dy=None, name=None):
+		mirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.R = R
+		self.CA = CA
+		self.type = 'Spherical mirror'
+
+	def phaseDelayFunction(self, x, y, l):
+		_k = 2 * np.pi / l
+		_dz = self.R * (1 - np.sqrt(1 - (x**2 + y**2)/self.R**2) )
+		_t_m = np.exp(-1j * _k * 2 * _dz)
+		return _t_m
+
+
+
+class freeformMirror(mirrorObject):
+	def __init__(self, x, y, z, dx=None, dy=None, name=None):
+		mirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.x = x
+		self.y = y
+		self.z = z
+		self.type = 'User def. mirror'
+
+	def phaseDelayFunction(self, l):
+		_k = 2 * np.pi / l
+		_t_m = np.exp(-1j * _k * 2 * self.z)
+		return _t_m
+
+
+
+
+class OpticalTrain:
+	'''
+		The main class where an optical system is build up sequentially.
+	'''
+	def __init__(self):
+		self.opticalTrain 		= []
+		self.sequentialDistance = []
+
+	def add_object(self, obj, dist):
+		'''
+			Adding a lens or mirror object to the opticalTrain
+			obj 	The lens or mirror object
+			dist 	The distance from the the previous object.
+		'''
+		assert type(dist) == float or type(dist) == int, 'Distance is not a float or integer.'
+		self.opticalTrain.append(obj)
+		self.sequentialDistance.append(float(dist))
+
+	def print_lens_Data(self):
+		'''
+			Print the data of the current opticalTrain.
+		'''
+		_tableData = []
+		for _lens in self.opticalTrain:
+			_i = self.opticalTrain.index(_lens)
+			_distances = self.sequentialDistance
+			_distances.append('-')
+			_tableData.append([_i, _lens.name, _lens.type, _distances[_i+1] ])
+		print(tabulate(_tableData, headers=['#', 'Name', 'Type', 'Distance to next surface']))
+
+
+	def propagate_system(self, U_in, x, y, l, output=None, paraxial=True):
+		# We need to have an input distribution
+		# We need to determine what propagation method is required
+		# Call the objects in the self.opticalTrain sequentially
+		# Call the subsequent phaseDelay
+		if output == 'complete':
+			_wf_list 	= []
+			_x_list 	= []
+			_y_list 	= []
+			_wf_list.append(_wf)
+			_x_list.append(x)
+			_y_list.append(y)
+		
+		_wf = U_in
+
+		for _lens in self.opticalTrain:
+			_i 	= self.opticalTrain.index(_lens)
+			_dz = self.sequentialDistance[_i]
+			_wf = _lens.phaseDelayFunction(x, y, l, paraxial=paraxial)
+
+			# Perform Fresnel number test:
+			_F = self.fresnel_number()
+
+			if _F < 5:
+				# Fresnel propagation
+				_wf, x, y = self.fresnel_propagation()
+			elif _F >= 5:
+				# Angular spectrum (p2p) propagation
+				_wf, x, y = self.angular_spectrum_propagation()
+
+			if output == 'complete':
+				_wf_list.append(_wf)
+				_x_list.append(x)
+				_y_list.append(y)
+		
+		if output == 'Complete':
+			return np.array([_wf_list, _x_list, _y_list])
+
+		else:
+			return np.array([_wf, x, y])
+
+		
+
+
+
+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.spiderTopHat(250, 100, 2001, 2001, dpix=1, oversample=1)
+	A = prof.topHat(2.5, n, m, dpix=dpix, oversample=1)
+
+	# Alternatively, create a gaussian profile
+	# gauss = GaussianBeam(0.3, l*1000)
+	# A = gauss.fundamentalMode(0,n,m,np.max(x))
+
+	# 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))
+	
+	# Start an optical train object. Add lenses and propagate. Currently not in use.
+	# LDE = OpticalTrain()
+
+	# 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.phaseDelayFunction(x,y,1.35e-6, 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.
+	oC = [-5, 5, -5, 5]
+	savename = 'simple_lens.pdf'
+	Exyz,_,_ = fou.FresnelPropagationFFTW(U_in, x, y, dz=L1.f, oversample=None, show=True, save=savename, outputCoordinates=oC)
+	
+	# 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 True:
+		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.FresnelPropagationFFTW(U_in, xin, yin, dz=i, oversample=None, 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(211)
+		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(212)
+		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)
+
+		# 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)))
+	
+
+

test cases (standalone)/fourierOptics_02_telescope.py

@@ -0,0 +1,771 @@
+#!/usr/bin/env Python3
+import numpy as np
+from pdb import set_trace
+import matplotlib.pyplot as plt
+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 time
+
+
+class FourierOptics:
+	def __init__(self, l):
+		self.l = l # Wavelength
+
+
+	def planeWavePropagationFFTW(self, A, dpix=None, dz=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.abs(_E_xy0)**2).sum()	# Factor to scale the intensity after the FFTs
+		_IntensityNormFactor = np.sum(np.abs(_E_xy0)**2)
+		_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)
+		_kz = self.phaseFactorExponent(_fx, _fy, self.l) 	# origin at start of array
+		_phaseFactor = np.exp(-1j * _kz * dz) 				# origin at start of array
+		
+		# Calculating the spatial frequencies in k-space
+		_kx 		= np.fft.fftshift( _fx * 2 * np.pi )
+		_ky 		= np.fft.fftshift( _fy * 2 * np.pi )
+
+		_E_xyz 		= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_A_in 		= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_A_in[:] 	= _Az0 * _phaseFactor
+
+		_FFTW_Obj2 	= pyfftw.FFTW( _A_in, _E_xyz, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
+		_FFTW_Obj2()
+		_E_xyz 		= np.fft.fftshift(_E_xyz) 				# use fftshift to move the origin from edge to center
+		_E_xyz /= np.sqrt(_IntensityNormFactor) 			# Normalization back to initial energy.
+
+		# Coordinates of E_xy0
+		_y, _x = np.mgrid[-_n / 2 * dpix : _n / 2 * dpix : 1j*_n, -_m / 2 * dpix : _m / 2 * dpix : 1j * _m]
+		if outputLimits:
+			# 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.
+			_x				= np.linspace(-_n/2*dpix, _n/2*dpix, _n)
+			_y				= np.linspace(-_m/2*dpix, _m/2*dpix, _m)
+			_f_ip_real		= interpolate.interp2d(_x, _y, _E_xyz.real, kind='linear')
+			_f_ip_imag 		= interpolate.interp2d(_x, _y, _E_xyz.imag, kind='linear')
+			_xi				= np.linspace(outputLimits[0],outputLimits[1],_n)
+			_yi 			= np.linspace(outputLimits[0],outputLimits[1],_m)
+			_E_out_ip_real 	= _f_ip_real(_xi, _yi)
+			_E_out_ip_imag	= _f_ip_imag(_xi, _yi)
+			_E_xyz 			= _E_out_ip_real + _E_out_ip_imag*1j
+			_y, _x 			= np.mgrid[outputLimits[0]:outputLimits[1]:1j*_n, outputLimits[0]:outputLimits[1]:1j * _m]
+			
+
+		_E_xyz_extent = [np.min(_x), np.max(_x), np.min(_y), np.max(_x)]
+
+
+
+		if show or save:
+			fig = plt.figure(figsize=(8,3))
+			ax = fig.add_subplot(121)
+			im = ax.imshow(np.abs(_E_xyz), cmap='hot', extent=_E_xyz_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_xyz), cmap='hot', extent=_E_xyz_extent)
+			ax2.set_xlabel('x [mm]')
+			ax2.set_ylabel('y [mm]')
+			fig.colorbar(im2, ax=ax2)
+			plt.tight_layout()
+		if save:
+			plt.savefig(save, bbox_inches='tight')
+		if show:
+			plt.show()
+		elif not show:
+			plt.close()
+
+		return np.array([_E_xyz, _x, _y])
+
+
+	def FresnelPropagationFFTW(self, A, xin, yin, dz=None, oversample=None, show=False, save=False, scaleInputArray=False, outputCoordinates='auto'):
+		'''
+			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]
+											If 'auto': the output coordinates will be determined 
+											automatically.
+
+			Output:
+			E_out 				- 	[array] (Complex) Output field distribution
+			x, y 				-	[array] Output coordinates 
+		'''
+		assert np.shape(xin) == np.shape(A), 'Input coordinates (xin) do not match shape of input field.'
+		assert np.shape(yin) == np.shape(A), 'Input coordinates (yin) do not match shape of input field.'
+		A = np.conj(A)
+
+		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='linear')
+			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)
+		_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
+		
+		# Calculate the pixel size from the xi, eta plane
+		_dpix = np.mean(np.diff(_xi))
+		if _dpix < 1e-99:
+			_dpix = np.mean(np.diff(_xi.T))
+
+		# Gaussian test: use 8*w0 radius for FFT:
+		if outputCoordinates == 'auto':
+			_w0_fit, _ 		= leastsq(self.minimize, self.l, args=(0, _xi, _eta, A))
+			_gauss 			= GaussianBeam(_w0_fit, self.l)
+			_halfwidth 		= 8 * _gauss.beamRadius(dz)[0]
+			_xip			= np.linspace(-_halfwidth, _halfwidth, _n) 						# Desired x,y - plane
+			_yip			= np.linspace(-_halfwidth, _halfwidth, _m) 						# Desired x,y - plane
+		else:
+			_xip			= np.linspace(outputCoordinates[0], outputCoordinates[1], _n) 	# Desired x,y - plane
+			_yip			= np.linspace(outputCoordinates[2], outputCoordinates[3], _m) 	# Desired x,y - plane
+
+
+		# The quadratic phase exponential of the input wavefront
+		_phaseComponent = np.exp(1j * _k / (2 * dz) * (_xi**2 + _eta**2))
+
+		# Complete input
+		_U_in 			= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_U_in[:]		= A * _phaseComponent
+		
+		# Output of the FFT evaluated at frequency fx = x / (lambda * z)
+		_U_out 			= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_FFTW_Obj		= pyfftw.FFTW( np.fft.ifftshift(_U_in), _U_out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
+		_FFTW_Obj()
+		_U_out 			= np.fft.fftshift(_U_out)
+
+		# Evaluate frequencies:
+		_fx 			= np.fft.fftshift(np.fft.fftfreq(np.shape(_xi)[0], d=_dpix))
+		_fy 			= np.fft.fftshift(np.fft.fftfreq(np.shape(_xi)[1], d=_dpix))
+		_x 				= _fx * self.l * dz
+		_y 				= _fy * self.l * dz
+		_x2, _y2 		= np.meshgrid(_x, _y)
+
+
+		# Final output and frequencies
+		_pre_amble 		= np.exp(1j * _k * dz) / (1j * self.l * dz) * np.exp(1j * _k / (2 * dz) * (_x2**2 + _y2**2))
+		_E_out 			= _pre_amble * _U_out
+
+		# Normalization of the output values
+		_IntensityNormFactor = np.sum(np.abs(A)**2)
+		_E_out /= (np.abs(np.exp(1j*_k*dz)/(1j*self.l*dz)) * np.sqrt(_IntensityNormFactor) * np.sqrt(_n*_m))
+
+		# Performing interpolation to get to the required final plane, done for real and imaginary part.
+		_f_ip_real		= interpolate.interp2d(_x, _y, _E_out.real, kind='linear')
+		_f_ip_imag 		= interpolate.interp2d(_x, _y, _E_out.imag, kind='linear')
+		_E_out_ip_real 	= _f_ip_real(_xip, _yip)
+		_E_out_ip_imag	= _f_ip_imag(_xip, _yip)
+		_E_out_ip 		= _E_out_ip_real + _E_out_ip_imag*1j
+		_E_out_ip 		/= np.sqrt(((_x2.max()/_xip.max())*(_y2.max()/_yip.max())))
+		_E_out_ip = np.conj(_E_out_ip)
+
+		if show:
+			# Plotting the results
+			_fig = plt.figure(figsize=(8,3))
+			_ax = _fig.add_subplot(121)
+			# _im = _ax.imshow(np.sqrt(np.abs(_E_out_ip)), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()])
+			_im = _ax.imshow(np.abs(_E_out_ip), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()])
+			_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='Amplitude')
+			# _fig.colorbar(_im, ax=_ax, label='$\\sqrt{Amplitude}$')
+			_ax2 = _fig.add_subplot(122)
+			_im2 = _ax2.imshow(np.angle(_E_out_ip), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()])
+			_ax2.set_xlim([_xip.min(), _xip.max()])
+			_ax2.set_ylim([_yip.min(), _yip.max()])
+			_ax2.set_xlabel('x [mm]')
+			_ax2.set_ylabel('y [mm]')
+			_fig.colorbar(_im2, ax=_ax2, label='Phase')
+			plt.tight_layout()
+			if save:
+				plt.savefig(save, bbox_inches='tight', transparent=True)
+			plt.show()
+
+		_x, _y = np.meshgrid(_xip, _yip)
+		return np.array([_E_out_ip, _x, _y])
+
+
+	def freqArray(self, A, dpix):
+		n, m = np.shape(A)
+		
+		fy, fx = np.mgrid[0:m, 0:n]
+
+		fx = fx.flatten()
+		fy = fy.flatten()
+		
+		for i in range(len(fx)):
+			fx[i] = (fx[i] - n*bool(fx[i]>(n-1)/2) ) / (n*dpix)
+			fy[i] = (fy[i] - m*bool(fy[i]>(m-1)/2) ) / (m*dpix)
+		
+		fx = fx.reshape(n,m)
+		fy = fy.reshape(n,m)
+		
+		return fx, fy
+				
+
+	def phaseFactorExponent(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 minimize(self, a, dz, x, y, og_data):
+		'''
+			Residual function used for fitting the gaussian beam to the input distribution.
+			Input:
+			a 		- 	[float] Guess for w0
+			dz  	- 	[float] Propagation distance, should be known
+			x,y 	- 	[array] meshgrid of the x and y coordinates.
+			og_data - 	[array] input distribution to be fitted.
+		'''
+		return np.abs((self.gaussianBeamFit(a,dz,x,y) - og_data).flatten()).real
+
+
+	def gaussianBeamFit(self, a, dz, x, y):
+		'''
+			Function used to fit the Gaussian beam to an input plane. Very similar
+			to the GaussianBeam class below, but modified to allow a best fit.
+			Input:
+			a 	- 	[float] Guess for w0
+			dz  - 	[float] Propagation distance, should be known
+			x,y - 	[array] meshgrid of the x and y coordinates.
+		'''
+		_w0 = a
+		_z 	= dz
+		_r 	  	= np.sqrt(x**2 + y**2)
+		if _z != 0:
+			_z_c 	= np.pi * _w0**2 / self.l
+			_w 		= _w0 * np.sqrt(1 + (_z/_z_c)**2)
+			_R 		= _z + _z_c **2 / _z
+			_phi0 	= np.arctan(_z/_z_c)
+			_k 		= 2 * np.pi / self.l
+			_results = np.sqrt(2 / (np.pi*_w**2)) * np.exp(- _r**2 / _w**2 - 1j * _k * _z - 1j * np.pi * _r**2 / (self.l * _R) + 1j*_phi0)
+		elif _z == 0:
+			_results = np.sqrt(2 / (np.pi*_w0**2)) * np.exp( - _r**2 / _w0**2 )
+		return _results
+
+
+
+
+
+
+class profiles:
+	def __init__(self):
+		pass
+
+	def topHat(self, r, n, m, dpix, oversample=1):
+	    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 <= (r)).astype(int)
+	    return _Anew
+
+	def obscuredTopHat(self, r1, r2, n, m, dpix, oversample=1):
+		'''
+			Creates a grid of n x m points, following a top hat function
+			where all points inside a radius r is 1 and 0 otherwise.
+		'''
+		assert r1 > r2, "r2 < r1, change input"
+		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))
+		return _Anew.astype(int)
+
+
+	def spiderTopHat(self, r1, r2, n, m, dpix, oversample=1):
+		'''
+			Creates a grid of n x m points, following a top hat function
+			where all points inside a radius r is 1 and 0 otherwise.
+		'''
+		assert r1 > r2, "r2 < r1, change input"
+		j,i 				= np.mgrid[-n/2*dpix : n/2*dpix : 1j*n*oversample, -m/2*dpix : m/2*dpix : 1j*m*oversample]
+		n *= oversample
+		m *= oversample
+		dist_from_center 	= np.sqrt(i**2 + j**2)
+		_Anew 				= (dist_from_center <= (r1))  & (dist_from_center >= (r2))
+		_Aspider 			= (np.abs(i) < (0.05*r1)) | (np.abs(j) < (0.05*r1))
+		_Anew 				= _Anew & ~_Aspider
+		return _Anew
+
+
+
+class GaussianBeam:
+	def __init__(self, w0, l):
+		# Confocal distance
+		self.z_c 	 = np.pi * w0**2 / l
+
+		# Far-field divergence angle
+		self.theta_0 = l / (np.pi * w0)
+		self.l 		 = l
+		self.w0 	 = w0
+		self.fwhm 	 = 1.18 * self.theta_0
+
+	def fundamentalMode(self, z, n, m, rmax):
+		y, x = np.mgrid[-rmax:rmax:n*1j,-rmax:rmax:n*1j]
+		r 	 = np.sqrt(x**2 + y**2)
+		if z != 0:
+			_w 		= self.beamRadius(z)
+			_R 		= self.curvature(z)
+			_phi0 	= self.phaseShift(z)
+			_k 		= 2 * np.pi / self.l
+			return np.sqrt(2 / (np.pi*_w**2)) * np.exp(- r**2 / _w**2 - 1j * _k * z - 1j * np.pi * r**2 / (self.l * _R) + 1j*_phi0)
+		elif z == 0:
+			return np.sqrt(2 / (np.pi*self.w0**2)) * np.exp( - r**2 / self.w0**2 ) 
+
+	def curvature(self, z):
+		return z + self.z_c **2 / z
+
+	def beamRadius(self, z):
+		return self.w0 * np.sqrt(1 + (z/self.z_c)**2)
+
+	def phaseShift(self, z):
+		return np.arctan(z/self.z_c)
+
+	def relativePower(self, r, z):
+		return np.exp(-2 * (r / self.beamRadius(z))**2)
+
+
+
+class lensObject:
+	'''
+		The lens object describes any spherical lens and find the phase delay
+		for all light traveling through it.
+		The refractive index, of course, depends on the wavelength.
+	'''
+	def __init__(self, glass, R1, R2, thickness, CA, f=None, name=None):
+		self.glass 		= glass
+		assert glassCatalog.catalog.get(self.glass), 'Glass type not in catalog'
+		self.R1 		= R1 				# mm
+		self.R2 		= R2 				# mm
+		self.thickness 	= thickness			# mm
+		self.CA 		= CA 				# mm
+		self.dispersionCoefficients = glassCatalog.catalog.get(self.glass)
+		self.referenceWavelength = 1.35 	# micron
+		self.type 		= 'Spherical lens'
+		if not f:
+			self.f 		= self.findFocalLength()
+		else:
+			self.f 		= f
+		if name:
+			assert type(name)==str, 'No string for a name.'
+			self.name = name
+		else:
+			self.name = str(self)
+
+	def findRefractiveIndex(self, _l):
+		_a = self.dispersionCoefficients
+		_n = np.sqrt( (_a[0]* _l**2)/(_l**2 - _a[1]) + (_a[2] * _l**2) / (_l**2 - _a[3]) + (_a[4] * _l**2) / (_l**2 - _a[5]) + 1)
+		return _n
+
+	def findFocalLength(self):
+		_n 		= self.findRefractiveIndex(self.referenceWavelength)
+		_f_recip = (_n - 1) * (1 / self.R1 - 1 / self.R2)
+		return 1/_f_recip
+
+	def phaseDelayFunction(self, x, y, l, paraxial=False):
+		if not paraxial:
+			_A1 = self.R1 * (1 - np.sqrt(1 - (x**2 + y**2) / self.R1**2 ) )
+			_A2 = self.R2 * (1 - np.sqrt(1 - (x**2 + y**2) / self.R2**2 ) )
+			_dxy = self.thickness - _A1 + _A2
+		if paraxial:
+			_dxy = self.thickness - (x**2 + y**2)/2 * (1/self.R1 - 1/self.R2)
+		_k = 2 * np.pi / 30e-3
+		_n = self.findRefractiveIndex(self.referenceWavelength)
+		_t_l = np.exp(-1j * _k * self.thickness)* np.exp(-1j * _k * (_n - 1) * _dxy)
+
+		# 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)).astype(int)
+		_t_l *= _lensMask
+		return _t_l
+
+
+
+class mirrorObject:
+	'''
+		Describes any mirror. No wavelength dependence exist. 
+		Coatings are not considered.
+	'''
+	def __init__(self, dx, dy, name=None):
+		self.dx = dx
+		self.dy = dy
+		if name:
+			assert type(name)==str, 'No string for a name.'
+			self.name = name
+		else:
+			self.name = str(self)
+
+
+class parabolicMirror(mirrorObject):
+	def __init__(self, f, CA, dx=None, dy=None, name=None):
+		mirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.f = f
+		self.CA = CA
+		self.type = 'Parabolic mirror'
+
+	def phaseDelayFunction(self, x, y, l):
+		_k 		= 2 * np.pi / l
+		_dz 	= (x**2 + y**2) / (4 * self.f)
+		_t_m 	= np.exp(1j * _k * 2 * _dz)
+		return _t_m
+
+
+
+class sphericalMirror(mirrorObject):
+	def __init__(self, R, CA, dx=None, dy=None, name=None):
+		mirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.R = R
+		self.CA = CA
+		self.type = 'Spherical mirror'
+
+	def phaseDelayFunction(self, x, y, l):
+		# _k = 2 * np.pi / l
+		# _dz = self.R * (1 - np.sqrt(1 - (x**2 + y**2)/self.R**2) )
+		# _t_m = np.exp(1j * _k * 2 * _dz)
+		# return _t_m
+		_k 		= 2 * np.pi / l
+		_dz 	= (x**2 + y**2) / (4 * self.f)
+		
+		# 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)
+		_t_m *= _mirrorMask
+		return _t_m
+
+
+class freeformMirror(mirrorObject):
+	def __init__(self, x, y, z, dx=None, dy=None, name=None):
+		mirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.x = x
+		self.y = y
+		self.z = z
+		self.type = 'User def. mirror'
+
+	def phaseDelayFunction(self, l):
+		_k = 2 * np.pi / l
+		_t_m = np.exp(1j * _k * 2 * self.z)
+		return _t_m
+
+
+
+
+class OpticalTrain:
+	'''
+		The main class where an optical system is build up sequentially.
+	'''
+	def __init__(self):
+		self.opticalTrain 		= []
+		self.sequentialDistance = []
+
+	def add_object(self, obj, dist):
+		'''
+			Adding a lens or mirror object to the opticalTrain
+			obj 	The lens or mirror object
+			dist 	The distance from the the previous object.
+		'''
+		assert type(dist) == float or type(dist) == int, 'Distance is not a float or integer.'
+		self.opticalTrain.append(obj)
+		self.sequentialDistance.append(float(dist))
+
+	def print_lens_Data(self):
+		'''
+			Print the data of the current opticalTrain.
+		'''
+		_tableData = []
+		for _lens in self.opticalTrain:
+			_i = self.opticalTrain.index(_lens)
+			_distances = self.sequentialDistance
+			_distances.append('-')
+			_tableData.append([_i, _lens.name, _lens.type, _distances[_i+1] ])
+		print(tabulate(_tableData, headers=['#', 'Name', 'Type', 'Distance to next surface']))
+
+
+	def propagate_system(self, U_in, x, y, l, output=None, paraxial=True):
+		# We need to have an input distribution
+		# We need to determine what propagation method is required
+		# Call the objects in the self.opticalTrain sequentially
+		# Call the subsequent phaseDelay
+		if output == 'complete':
+			_wf_list 	= []
+			_x_list 	= []
+			_y_list 	= []
+			_wf_list.append(_wf)
+			_x_list.append(x)
+			_y_list.append(y)
+		
+		_wf = U_in
+
+		for _lens in self.opticalTrain:
+			_i 	= self.opticalTrain.index(_lens)
+			_dz = self.sequentialDistance[_i]
+			_wf = _lens.phaseDelayFunction(x, y, l, paraxial=paraxial)
+
+			# Perform Fresnel number test:
+			_F = self.fresnel_number()
+
+			if _F < 5:
+				# Fresnel propagation
+				_wf, x, y = self.fresnel_propagation()
+			elif _F >= 5:
+				# Angular spectrum (p2p) propagation
+				_wf, x, y = self.angular_spectrum_propagation()
+
+			if output == 'complete':
+				_wf_list.append(_wf)
+				_x_list.append(x)
+				_y_list.append(y)
+		
+		if output == 'Complete':
+			return np.array([_wf_list, _x_list, _y_list])
+
+		else:
+			return np.array([_wf, x, y])
+
+		
+
+
+
+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 = 8
+	n, m = 4001, 4001
+	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.topHat(4000, n, m, dpix=dpix, oversample=1)
+	# A = prof.spiderTopHat(4000, 1000, n, m, dpix=dpix, oversample=1)
+	# A = prof.obscuredTopHat(4000, 1000, n, m, dpix=dpix, oversample=1)
+
+	# Alternatively, create a gaussian profile
+	# gauss = GaussianBeam(0.3, l*1000)
+	# A = gauss.fundamentalMode(0,n,m,np.max(x))
+
+	# 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))
+	
+	# Start an optical train object. Add lenses and propagate. Currently not in use.
+	# LDE = OpticalTrain()
+
+	# Create a mirror with a focal length 64 m, with a clear aperture radius of 4 m.
+	M1 = parabolicMirror(dx=0, dy=0, name='M1', CA=8000, f=64000)
+	print('Focal length of mirror {} is {:.1f} mm.'.format(M1.name, M1.f))
+
+	# Calculate the phase delay introduced by the mirror
+	t_l = M1.phaseDelayFunction(x,y,l)
+	
+	# The input distribution gets a curved wavefront after the mirror
+	U_in = A * t_l
+	
+	plt.imshow(np.abs(U_in))
+	plt.colorbar()
+	plt.savefig('telescope_mask.pdf', bbox_inches='tight', transparent=True)
+	plt.show(block=False)
+
+	# Performing the propagation using Fresnel approximation to the focal plane.
+	# oC = [-0.001, 0.001, -0.001, 0.001]
+	oC = [-2, 2, -2, 2]
+	# savename = 'fourierOptics_02_telescope_fresnel.pdf'
+	Exyz,xout,yout = fou.FresnelPropagationFFTW(U_in, x, y, dz=M1.f, oversample=None, show=True, save=False, outputCoordinates=oC)
+
+	# 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[int((n-1)/2),:]**2), 'r-')
+	ax.set_xlabel('x (mm)')
+	ax.set_ylabel('Intensity')
+	ax.tick_params(axis ='both', direction='in', which='both') 
+	# plt.savefig('fourierOptics_02_telescope_slice.pdf', bbox_inches='tight', transparent=True)
+	plt.show()
+
+
+	# Combined plot
+	# Plotting the results
+	_fig = plt.figure(figsize=(12,3))
+	_ax = _fig.add_subplot(131)
+	_im = _ax.imshow(np.abs(Exyz), cmap='hot', extent=oC)
+	_ax.set_xlim([oC[0], oC[1]])
+	_ax.set_ylim([oC[2], oC[3]])
+	_ax.set_xlabel('x [mm]')
+	_ax.set_ylabel('y [mm]')
+	_fig.colorbar(_im, ax=_ax, label='Amplitude')
+	_ax2 = _fig.add_subplot(132)
+	_im2 = _ax2.imshow(np.angle(Exyz), cmap='hot', extent=oC)
+	_ax2.set_xlim([oC[0], oC[1]])
+	_ax2.set_ylim([oC[2], oC[3]])
+	_ax2.set_xlabel('x [mm]')
+	_ax2.set_ylabel('y [mm]')
+	_fig.colorbar(_im2, ax=_ax2, label='Phase')
+	_ax3 = _fig.add_subplot(133)
+	_ax3.grid(True, alpha=0.3)
+	_ax3.semilogy(np.linspace(np.min(oC), np.max(oC),n), np.abs(Exyz[int((n-1)/2),:]**2), 'r-')
+	_ax3.set_xlabel('x (mm)')
+	_ax3.set_ylabel('Intensity')
+	_ax3.tick_params(axis ='both', direction='in', which='both') 
+	plt.tight_layout()
+	plt.savefig('fourierOptics_02_telescope_complete.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.FresnelPropagationFFTW(U_in, xin, yin, dz=i, oversample=None, 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(211)
+		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(212)
+		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('Intensity')
+		ax2.set_xlim([np.min(zList), np.max(zList)])
+		ax2.grid(True, alpha=0.3)
+
+		# 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.show()
+
+
+	# Closing down the script and stopping the timer.
+	end = time.time()
+	print('The program took {:.3f} seconds!'.format((end-start)))
+	
+
+

test cases (standalone)/fourierOptics_03_verification.py

@@ -0,0 +1,775 @@
+#!/usr/bin/env Python3
+import numpy as np
+from pdb import set_trace
+import matplotlib.pyplot as plt
+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
+
+class FourierOptics:
+	def __init__(self, l):
+		self.l = l # Wavelength
+
+
+	def planeWavePropagationFFTW(self, A, dpix=None, dz=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.abs(_E_xy0)**2).sum()	# Factor to scale the intensity after the FFTs
+		_IntensityNormFactor = np.sum(np.abs(_E_xy0)**2)
+		_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)
+		_kz = self.phaseFactorExponent(_fx, _fy, self.l) 	# origin at start of array
+		_phaseFactor = np.exp(-1j * _kz * dz) 				# origin at start of array
+		
+		# Calculating the spatial frequencies in k-space
+		_kx 		= np.fft.fftshift( _fx * 2 * np.pi )
+		_ky 		= np.fft.fftshift( _fy * 2 * np.pi )
+
+		_E_xyz 		= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_A_in 		= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_A_in[:] 	= _Az0 * _phaseFactor
+
+		_FFTW_Obj2 	= pyfftw.FFTW( _A_in, _E_xyz, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
+		_FFTW_Obj2()
+		_E_xyz 		= np.fft.fftshift(_E_xyz) 				# use fftshift to move the origin from edge to center
+		_E_xyz /= np.sqrt(_IntensityNormFactor) 			# Normalization back to initial energy.
+
+		# Coordinates of E_xy0
+		_y, _x = np.mgrid[-_n / 2 * dpix : _n / 2 * dpix : 1j*_n, -_m / 2 * dpix : _m / 2 * dpix : 1j * _m]
+		if outputLimits:
+			# 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.
+			_x				= np.linspace(-_n/2*dpix, _n/2*dpix, _n)
+			_y				= np.linspace(-_m/2*dpix, _m/2*dpix, _m)
+			_f_ip_real		= interpolate.interp2d(_x, _y, _E_xyz.real, kind='linear')
+			_f_ip_imag 		= interpolate.interp2d(_x, _y, _E_xyz.imag, kind='linear')
+			_xi				= np.linspace(outputLimits[0],outputLimits[1],_n)
+			_yi 			= np.linspace(outputLimits[0],outputLimits[1],_m)
+			_E_out_ip_real 	= _f_ip_real(_xi, _yi)
+			_E_out_ip_imag	= _f_ip_imag(_xi, _yi)
+			_E_xyz 			= _E_out_ip_real + _E_out_ip_imag*1j
+			_y, _x 			= np.mgrid[outputLimits[0]:outputLimits[1]:1j*_n, outputLimits[0]:outputLimits[1]:1j * _m]
+			
+
+		_E_xyz_extent = [np.min(_x), np.max(_x), np.min(_y), np.max(_x)]
+
+
+
+		if show or save:
+			fig = plt.figure(figsize=(8,3))
+			ax = fig.add_subplot(121)
+			im = ax.imshow(np.abs(_E_xyz), cmap='hot', extent=_E_xyz_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_xyz), cmap='hot', extent=_E_xyz_extent)
+			ax2.set_xlabel('x [mm]')
+			ax2.set_ylabel('y [mm]')
+			fig.colorbar(im2, ax=ax2)
+			plt.tight_layout()
+		if save:
+			plt.savefig(save, bbox_inches='tight')
+		if show:
+			plt.show()
+		elif not show:
+			plt.close()
+
+		return np.array([_E_xyz, _x, _y])
+
+
+	def FresnelPropagationFFTW(self, A, xin, yin, dz=None, oversample=None, show=False, save=False, scaleInputArray=False, outputCoordinates='auto'):
+		'''
+			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]
+											If 'auto': the output coordinates will be determined 
+											automatically.
+
+			Output:
+			E_out 				- 	[array] (Complex) Output field distribution
+			x, y 				-	[array] Output coordinates 
+		'''
+		assert np.shape(xin) == np.shape(A), 'Input coordinates (xin) do not match shape of input field.'
+		assert np.shape(yin) == np.shape(A), 'Input coordinates (yin) do not match shape of input field.'
+		A = np.conj(A)
+
+		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='linear')
+			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)
+		_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
+		
+		# Calculate the pixel size from the xi, eta plane
+		_dpix = np.mean(np.diff(_xi))
+		if _dpix < 1e-99:
+			_dpix = np.mean(np.diff(_xi.T))
+
+		# Gaussian test: use 8*w0 radius for FFT:
+		if outputCoordinates == 'auto':
+			_w0_fit, _ 		= leastsq(self.minimize, self.l, args=(0, _xi, _eta, A))
+			_gauss 			= GaussianBeam(_w0_fit, self.l)
+			_halfwidth 		= 8 * _gauss.beamRadius(dz)[0]
+			_xip			= np.linspace(-_halfwidth, _halfwidth, _n) 						# Desired x,y - plane
+			_yip			= np.linspace(-_halfwidth, _halfwidth, _m) 						# Desired x,y - plane
+		else:
+			_xip			= np.linspace(outputCoordinates[0], outputCoordinates[1], _n) 	# Desired x,y - plane
+			_yip			= np.linspace(outputCoordinates[2], outputCoordinates[3], _m) 	# Desired x,y - plane
+
+
+		# The quadratic phase exponential of the input wavefront
+		_phaseComponent = np.exp(1j * _k / (2 * dz) * (_xi**2 + _eta**2))
+
+		# Complete input
+		_U_in 			= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_U_in[:]		= A * _phaseComponent
+		
+		# Output of the FFT evaluated at frequency fx = x / (lambda * z)
+		_U_out 			= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_FFTW_Obj		= pyfftw.FFTW( np.fft.ifftshift(_U_in), _U_out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
+		_FFTW_Obj()
+		_U_out 			= np.fft.fftshift(_U_out)
+
+		# Evaluate frequencies:
+		_fx 			= np.fft.fftshift(np.fft.fftfreq(np.shape(_xi)[0], d=_dpix))
+		_fy 			= np.fft.fftshift(np.fft.fftfreq(np.shape(_xi)[1], d=_dpix))
+		_x 				= _fx * self.l * dz
+		_y 				= _fy * self.l * dz
+		_x2, _y2 		= np.meshgrid(_x, _y)
+
+
+		# Final output and frequencies
+		_pre_amble 		= np.exp(1j * _k * dz) / (1j * self.l * dz) * np.exp(1j * _k / (2 * dz) * (_x2**2 + _y2**2))
+		_E_out 			= _pre_amble * _U_out
+
+		# Normalization of the output values
+		_IntensityNormFactor = np.sum(np.abs(A)**2)
+		_E_out /= (np.abs(np.exp(1j*_k*dz)/(1j*self.l*dz)) * np.sqrt(_IntensityNormFactor) * np.sqrt(_n*_m))
+
+		# Performing interpolation to get to the required final plane, done for real and imaginary part.
+		_f_ip_real		= interpolate.interp2d(_x, _y, _E_out.real, kind='linear')
+		_f_ip_imag 		= interpolate.interp2d(_x, _y, _E_out.imag, kind='linear')
+		_E_out_ip_real 	= _f_ip_real(_xip, _yip)
+		_E_out_ip_imag	= _f_ip_imag(_xip, _yip)
+		_E_out_ip 		= _E_out_ip_real + _E_out_ip_imag*1j
+		_E_out_ip 		/= np.sqrt(((_x2.max()/_xip.max())*(_y2.max()/_yip.max())))
+		_E_out_ip = np.conj(_E_out_ip)
+
+		if show:
+			# Plotting the results
+			_fig = plt.figure(figsize=(8,3))
+			_ax = _fig.add_subplot(121)
+			_im = _ax.imshow(np.abs(_E_out_ip), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()])
+			_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='Amplitude')
+			_ax2 = _fig.add_subplot(122)
+			_im2 = _ax2.imshow(np.angle(_E_out_ip), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()])
+			_ax2.set_xlim([_xip.min(), _xip.max()])
+			_ax2.set_ylim([_yip.min(), _yip.max()])
+			_ax2.set_xlabel('x [mm]')
+			_ax2.set_ylabel('y [mm]')
+			_fig.colorbar(_im2, ax=_ax2, label='Phase')
+			plt.tight_layout()
+			if save:
+				plt.savefig(save, bbox_inches='tight', transparent=True)
+			plt.show()
+
+		_x, _y = np.meshgrid(_xip, _yip)
+		return np.array([_E_out_ip, _x, _y])
+
+
+	def freqArray(self, A, dpix):
+		n, m = np.shape(A)
+		
+		fy, fx = np.mgrid[0:m, 0:n]
+
+		fx = fx.flatten()
+		fy = fy.flatten()
+		
+		for i in range(len(fx)):
+			fx[i] = (fx[i] - n*bool(fx[i]>(n-1)/2) ) / (n*dpix)
+			fy[i] = (fy[i] - m*bool(fy[i]>(m-1)/2) ) / (m*dpix)
+		
+		fx = fx.reshape(n,m)
+		fy = fy.reshape(n,m)
+		
+		return fx, fy
+				
+
+	def phaseFactorExponent(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 minimize(self, a, dz, x, y, og_data):
+		'''
+			Residual function used for fitting the gaussian beam to the input distribution.
+			Input:
+			a 		- 	[float] Guess for w0
+			dz  	- 	[float] Propagation distance, should be known
+			x,y 	- 	[array] meshgrid of the x and y coordinates.
+			og_data - 	[array] input distribution to be fitted.
+		'''
+		return np.abs((self.gaussianBeamFit(a,dz,x,y) - og_data).flatten()).real
+
+
+	def gaussianBeamFit(self, a, dz, x, y):
+		'''
+			Function used to fit the Gaussian beam to an input plane. Very similar
+			to the GaussianBeam class below, but modified to allow a best fit.
+			Input:
+			a 	- 	[float] Guess for w0
+			dz  - 	[float] Propagation distance, should be known
+			x,y - 	[array] meshgrid of the x and y coordinates.
+		'''
+		_w0 = a
+		_z 	= dz
+		_r 	  	= np.sqrt(x**2 + y**2)
+		if _z != 0:
+			_z_c 	= np.pi * _w0**2 / self.l
+			_w 		= _w0 * np.sqrt(1 + (_z/_z_c)**2)
+			_R 		= _z + _z_c **2 / _z
+			_phi0 	= np.arctan(_z/_z_c)
+			_k 		= 2 * np.pi / self.l
+			_results = np.sqrt(2 / (np.pi*_w**2)) * np.exp(- _r**2 / _w**2 - 1j * _k * _z - 1j * np.pi * _r**2 / (self.l * _R) + 1j*_phi0)
+		elif _z == 0:
+			_results = np.sqrt(2 / (np.pi*_w0**2)) * np.exp( - _r**2 / _w0**2 )
+		return _results
+
+
+
+
+
+
+class profiles:
+	def __init__(self):
+		pass
+
+	def topHat(self, r, n, m, dpix, oversample=1):
+	    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 <= (r)).astype(int)
+	    return _Anew
+
+	def obscuredTopHat(self, r1, r2, n, m, dpix, oversample=1):
+		'''
+			Creates a grid of n x m points, following a top hat function
+			where all points inside a radius r is 1 and 0 otherwise.
+		'''
+		assert r1 > r2, "r2 < r1, change input"
+		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))
+		return _Anew.astype(int)
+
+
+	def spiderTopHat(self, r1, r2, n, m, dpix, oversample=1):
+		'''
+			Creates a grid of n x m points, following a top hat function
+			where all points inside a radius r is 1 and 0 otherwise.
+		'''
+		assert r1 > r2, "r2 < r1, change input"
+		j,i 				= np.mgrid[-n/2*dpix : n/2*dpix : 1j*n*oversample, -m/2*dpix : m/2*dpix : 1j*m*oversample]
+		n *= oversample
+		m *= oversample
+		dist_from_center 	= np.sqrt(i**2 + j**2)
+		_Anew 				= (dist_from_center <= (r1))  & (dist_from_center >= (r2))
+		_Aspider 			= (np.abs(i) < (0.05*r1)) | (np.abs(j) < (0.05*r1))
+		_Anew 				= _Anew & ~_Aspider
+		return _Anew
+
+
+
+class GaussianBeam:
+	def __init__(self, w0, l):
+		# Confocal distance
+		self.z_c 	 = np.pi * w0**2 / l
+
+		# Far-field divergence angle
+		self.theta_0 = l / (np.pi * w0)
+		self.l 		 = l
+		self.w0 	 = w0
+		self.fwhm 	 = 1.18 * self.theta_0
+
+	def fundamentalMode(self, z, n, m, rmax):
+		y, x = np.mgrid[-rmax:rmax:n*1j,-rmax:rmax:n*1j]
+		r 	 = np.sqrt(x**2 + y**2)
+		if z != 0:
+			_w 		= self.beamRadius(z)
+			_R 		= self.curvature(z)
+			_phi0 	= self.phaseShift(z)
+			_k 		= 2 * np.pi / self.l
+			return np.sqrt(2 / (np.pi*_w**2)) * np.exp(- r**2 / _w**2 - 1j * _k * z - 1j * np.pi * r**2 / (self.l * _R) + 1j*_phi0)
+		elif z == 0:
+			return np.sqrt(2 / (np.pi*self.w0**2)) * np.exp( - r**2 / self.w0**2 ) 
+
+	def curvature(self, z):
+		return z + self.z_c **2 / z
+
+	def beamRadius(self, z):
+		return self.w0 * np.sqrt(1 + (z/self.z_c)**2)
+
+	def phaseShift(self, z):
+		return np.arctan(z/self.z_c)
+
+	def relativePower(self, r, z):
+		return np.exp(-2 * (r / self.beamRadius(z))**2)
+
+
+
+class lensObject:
+	'''
+		The lens object describes any spherical lens and find the phase delay
+		for all light traveling through it.
+		The refractive index, of course, depends on the wavelength.
+	'''
+	def __init__(self, glass, R1, R2, thickness, CA, f=None, name=None):
+		self.glass 		= glass
+		assert glassCatalog.catalog.get(self.glass), 'Glass type not in catalog'
+		self.R1 		= R1 				# mm
+		self.R2 		= R2 				# mm
+		self.thickness 	= thickness			# mm
+		self.CA 		= CA 				# mm
+		self.dispersionCoefficients = glassCatalog.catalog.get(self.glass)
+		self.referenceWavelength = 1.35 	# micron
+		self.type 		= 'Spherical lens'
+		if not f:
+			self.f 		= self.findFocalLength()
+		else:
+			self.f 		= f
+		if name:
+			assert type(name)==str, 'No string for a name.'
+			self.name = name
+		else:
+			self.name = str(self)
+
+	def findRefractiveIndex(self, _l):
+		_a = self.dispersionCoefficients
+		_n = np.sqrt( (_a[0]* _l**2)/(_l**2 - _a[1]) + (_a[2] * _l**2) / (_l**2 - _a[3]) + (_a[4] * _l**2) / (_l**2 - _a[5]) + 1)
+		return _n
+
+	def findFocalLength(self):
+		_n 		= self.findRefractiveIndex(self.referenceWavelength)
+		_f_recip = (_n - 1) * (1 / self.R1 - 1 / self.R2)
+		return 1/_f_recip
+
+	def phaseDelayFunction(self, x, y, l, paraxial=False):
+		if not paraxial:
+			_A1 = self.R1 * (1 - np.sqrt(1 - (x**2 + y**2) / self.R1**2 ) )
+			_A2 = self.R2 * (1 - np.sqrt(1 - (x**2 + y**2) / self.R2**2 ) )
+			_dxy = self.thickness - _A1 + _A2
+		if paraxial:
+			_dxy = self.thickness - (x**2 + y**2)/2 * (1/self.R1 - 1/self.R2)
+		_k = 2 * np.pi / 30e-3
+		_n = self.findRefractiveIndex(self.referenceWavelength)
+		_t_l = np.exp(-1j * _k * self.thickness)* np.exp(-1j * _k * (_n - 1) * _dxy)
+
+		# 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)).astype(int)
+		_t_l *= _lensMask
+		return _t_l
+
+
+
+class mirrorObject:
+	'''
+		Describes any mirror. No wavelength dependence exist. 
+		Coatings are not considered.
+	'''
+	def __init__(self, dx, dy, name=None):
+		self.dx = dx
+		self.dy = dy
+		if name:
+			assert type(name)==str, 'No string for a name.'
+			self.name = name
+		else:
+			self.name = str(self)
+
+
+class parabolicMirror(mirrorObject):
+	def __init__(self, f, CA, dx=None, dy=None, name=None):
+		mirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.f = f
+		self.CA = CA
+		self.type = 'Parabolic mirror'
+
+	def phaseDelayFunction(self, x, y, l):
+		_k 		= 2 * np.pi / l
+		_dz 	= (x**2 + y**2) / (4 * self.f)
+		_t_m 	= np.exp(-1j * _k * 2 * _dz)
+		return _t_m
+
+
+
+class sphericalMirror(mirrorObject):
+	def __init__(self, R, CA, dx=None, dy=None, name=None):
+		mirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.R = R
+		self.CA = CA
+		self.type = 'Spherical mirror'
+
+	def phaseDelayFunction(self, x, y, l):
+		_k = 2 * np.pi / l
+		_dz = self.R * (1 - np.sqrt(1 - (x**2 + y**2)/self.R**2) )
+		_t_m = np.exp(-1j * _k * 2 * _dz)
+		return _t_m
+
+
+
+class freeformMirror(mirrorObject):
+	def __init__(self, x, y, z, dx=None, dy=None, name=None):
+		mirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.x = x
+		self.y = y
+		self.z = z
+		self.type = 'User def. mirror'
+
+	def phaseDelayFunction(self, l):
+		_k = 2 * np.pi / l
+		_t_m = np.exp(-1j * _k * 2 * self.z)
+		return _t_m
+
+
+
+
+class OpticalTrain:
+	'''
+		The main class where an optical system is build up sequentially.
+	'''
+	def __init__(self):
+		self.opticalTrain 		= []
+		self.sequentialDistance = []
+
+	def add_object(self, obj, dist):
+		'''
+			Adding a lens or mirror object to the opticalTrain
+			obj 	The lens or mirror object
+			dist 	The distance from the the previous object.
+		'''
+		assert type(dist) == float or type(dist) == int, 'Distance is not a float or integer.'
+		self.opticalTrain.append(obj)
+		self.sequentialDistance.append(float(dist))
+
+	def print_lens_Data(self):
+		'''
+			Print the data of the current opticalTrain.
+		'''
+		_tableData = []
+		for _lens in self.opticalTrain:
+			_i = self.opticalTrain.index(_lens)
+			_distances = self.sequentialDistance
+			_distances.append('-')
+			_tableData.append([_i, _lens.name, _lens.type, _distances[_i+1] ])
+		print(tabulate(_tableData, headers=['#', 'Name', 'Type', 'Distance to next surface']))
+
+
+	def propagate_system(self, U_in, x, y, l, output=None, paraxial=True):
+		# We need to have an input distribution
+		# We need to determine what propagation method is required
+		# Call the objects in the self.opticalTrain sequentially
+		# Call the subsequent phaseDelay
+		if output == 'complete':
+			_wf_list 	= []
+			_x_list 	= []
+			_y_list 	= []
+			_wf_list.append(_wf)
+			_x_list.append(x)
+			_y_list.append(y)
+		
+		_wf = U_in
+
+		for _lens in self.opticalTrain:
+			_i 	= self.opticalTrain.index(_lens)
+			_dz = self.sequentialDistance[_i]
+			_wf = _lens.phaseDelayFunction(x, y, l, paraxial=paraxial)
+
+			# Perform Fresnel number test:
+			_F = self.fresnel_number()
+
+			if _F < 5:
+				# Fresnel propagation
+				_wf, x, y = self.fresnel_propagation()
+			elif _F >= 5:
+				# Angular spectrum (p2p) propagation
+				_wf, x, y = self.angular_spectrum_propagation()
+
+			if output == 'complete':
+				_wf_list.append(_wf)
+				_x_list.append(x)
+				_y_list.append(y)
+		
+		if output == 'Complete':
+			return np.array([_wf_list, _x_list, _y_list])
+
+		else:
+			return np.array([_wf, x, y])
+
+
+def FGB_1D(w0,l,z,x):
+    k    = 2 * np.pi / l
+    z    = z + 1e-16
+    w    = w0 * np.sqrt( 1 + ( (l*z) / (np.pi*w0**2) )**2)
+    R    = z + 1/z * ( np.pi * w0**2 / l)**2
+    phi  = np.arctan2( (l * z) , (np.pi * w0**2) )
+
+    norm = (2./ np.pi)**(.25) * (1. / w0**2)**(.5)
+    beam = np.exp( -x**2/w**2 - (1j * k * z) - ((1j*np.pi*x**2)/(l*R)) + 1j*phi/2)
+
+    return beam * norm
+
+if __name__ == '__main__':
+	plt.rc('font', family='serif')
+	plt.rc('axes', labelsize='large')
+
+	import time
+	start = time.time()
+	l = 30e-3
+	# l = 1e-6
+	fou = FourierOptics(l=l)
+	dpix = l/6
+	n = int(500 / dpix * l)
+	m = int(500 / dpix * l)
+	n, m = 3001, 3001
+	print (n,m)
+	prof = profiles()
+	# A = prof.topHat(4*l, 2001, 2001, 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))
+
+	w0_init 	= 4*l
+	z 			= 400*l
+	wg = w0_init * np.sqrt( 1 + ( (l*z) / (np.pi*w0_init**2) )**2)
+	# rmax_init 	= 4 * w0_init 
+	# rmax_init = 1000/8*l
+
+	xGauss = np.linspace(-(n/2)*dpix, (n/2)*dpix, n)
+	rmax_init = np.max(xGauss)
+	# A 			= GaussianBeam(w0_init, l).fundamentalMode(z, n, m, rmax_init)
+	A0 		= FGB_1D(w0_init, l, 0, xGauss) 
+	A0 		/= np.sqrt(np.sum(np.abs(A0)**2))
+	A 		= FGB_1D(wg, l, z, xGauss)
+
+	A 		/= np.sqrt(np.sum(np.abs(A)**2))
+
+	A2D 	= GaussianBeam(w0_init, l).fundamentalMode(z, n, m, rmax_init)
+	A2D 	/= np.sqrt(np.sum(np.abs(A2D[int(n/2)+1,:])**2))
+
+	A2D0 	= GaussianBeam(w0_init, l).fundamentalMode(0, n, m, rmax_init)
+
+	Ap2p,xp2p,_ 	= fou.planeWavePropagationFFTW(A2D0, dz=z, dpix=dpix, show=False, save=False)
+	Ap2p /= np.sqrt(np.sum(np.abs(Ap2p[int(n/2),:])**2))
+
+	oC = [xGauss.min(), xGauss.max(), xGauss.min(), xGauss.max()]
+	
+	# savename= 'FourierOptics_03_verification_fresnel.pdf'
+	savename=None
+	Afren,xfren,_ 	= fou.FresnelPropagationFFTW(A2D0, x,y, dz=z, show=True, save=savename, outputCoordinates=oC)
+	Afren /= np.sqrt(np.sum(np.abs(Afren[int(n/2),:])**2))
+	# Afren = np.conj(Afren)
+
+	norm_fac = np.max(np.abs(A2D)[int(n/2),:])
+
+	fig = plt.figure()
+	ax = fig.add_subplot(111)
+	# ax.plot(np.abs(A0), label='z=0, analytical')
+	# ax.plot(np.abs(A)/norm_fac, label='z=%i, analytical' %(int(z/l)))
+	ax.plot(np.abs(A2D)[int(n/2),:]/norm_fac, 'r--', label='z=%i$\lambda$, analytical' %(int(z/l)))
+	ax.plot(np.abs(Ap2p)[int(n/2),:]/norm_fac, label='z=%i$\lambda$, plane to plane' %(int(z/l)))
+	ax.plot(np.abs(Afren)[int(n/2),:]/norm_fac, label='z=%i$\lambda$, Fresnel' %(int(z/l)))
+	ax.grid(alpha=0.3)
+	ax.set_xlabel('Location [pixel no.]')
+	ax.set_ylabel('Normalized amplitude')
+	ax.legend()
+	plt.savefig('FourierOptics_03_verification_amplitude.pdf', bbox_inches='tight', transparent=True)
+	
+	fig2 = plt.figure(figsize=(12,3))
+	# ax1 = fig2.add_subplot(151)
+	# ax1.plot(np.angle(A0), label='z=0, analytical')
+	# ax1.set_title('z=0, analytical')
+	# ax2 = fig2.add_subplot(131)
+	# ax2.plot(np.angle(A), label='z=%i, analytical' %(int(z/l)))
+	# ax2.set_title('z=%i, analytical' %(int(z/l)))
+	ax3 = fig2.add_subplot(131)
+	ax3.plot(np.angle(A2D)[int(n/2),:], label='z=%i$\lambda$, analytical' %(int(z/l)))
+	ax3.set_title('z=%i$\lambda$, Gaussian beam' %(int(z/l)))
+	ax3.set_ylabel('Phase')
+	ax3.set_xlabel('Location [pixel no.]')
+	ax3.grid(True, alpha=0.3)
+	ax4 = fig2.add_subplot(132)
+	ax4.plot(np.angle(Ap2p)[int(n/2),:], label='z=%i$\lambda$, plane to plane' %(int(z/l)))
+	ax4.set_title('z=%i$\lambda$, Plane to plane' %(int(z/l)))
+	ax4.set_ylabel('Phase')
+	ax4.set_xlabel('Location [pixel no.]')
+	ax4.grid(True, alpha=0.3)
+	ax5 = fig2.add_subplot(133)
+	ax5.plot(np.angle(Afren)[int(n/2),:], label='z=%i$\lambda$, Fresnel' %(int(z/l)))
+	ax5.set_title('z=%i$\lambda$, Fresnel' %(int(z/l)))
+	ax5.set_ylabel('Phase')
+	ax5.set_xlabel('Location [pixel no.]')
+	ax5.grid(True, alpha=0.3)
+	# ax1.grid(alpha=0.3)
+	# ax6 = fig2.add_subplot(166)
+	# ax6.legend(ax=ax6)
+	plt.savefig('FourierOptics_03_verification_phase.pdf', bbox_inches='tight', transparent=True)
+	plt.tight_layout()
+	plt.show()
+
+
+	# set_trace()
+	# Afren,_,_ 	= fou.FresnelPropagationFFTW(A2D, dz=400*l, dpix=dpix, oversample=1, show=True, save=False, outputCoordinates=oC)
+
+	gauss = GaussianBeam(w0_init, l)
+	# A = prof.topHat(4*l, 2001, 2001, dpix=dpix, oversample=1)
+	A = gauss.fundamentalMode(0,n,m,rmax_init)
+	if False:
+
+		Ics 	= []
+		Ics_x 	= []
+		Isum 	= []
+		lList = np.linspace(1,201,10).tolist()
+		for i in lList:
+			title = f"{lList.index(i):03d}"
+			print (' Calculating {0}/{1}'.format(title, f"{len(lList):03d}"), end="\r")
+			# Exyz, x, y = fou.planeWavePropagationFFTW(A, dz=i*l, dpix=dpix, show=False, save=False)#, scaleInputArray=(n,m))
+			Exyz, x, y = fou.FresnelPropagationFFTW(A, dz=i*l, dpix=dpix, show=False, save=False, outputCoordinates=oC)
+			# Exyz = gauss.fundamentalMode(i*l, n, m, rmax_init)
+			# y, x = np.mgrid[-rmax_init:rmax_init:n*1j,-rmax_init:rmax_init:n*1j]
+			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]]
+			print(y_ext)
+			Exyz = None
+			x = None
+			y = None
+		
+		fig = plt.figure(figsize=(8,4))
+		ax  = fig.add_subplot(211)
+		ax.tick_params(axis ='both', direction='in') 
+		im = ax.imshow(np.asarray(Ics).T, cmap='hot', aspect='auto', interpolation='bicubic', extent=[lList[0],lList[-1],y_ext[0],y_ext[1]])
+		ax.yaxis.set_major_formatter(FormatStrFormatter('%.1E'))
+		# ax.set_ylim([-2.5e-4,2.5e-4])
+		ax2 = fig.add_subplot(212)
+		ax2.plot(lList, np.asarray(Ics)[:,int(np.asarray(Ics).shape[1]/2)], 'k-')
+		ax2.yaxis.set_major_formatter(FormatStrFormatter('%.1E'))
+		ax2.set_xlabel('Propagation distance [$\lambda$]')
+		ax2.tick_params(axis ='both', direction='in') 
+		ax2.set_ylabel('Intensity')
+		ax2.set_xlim([np.min(lList), np.max(lList)])
+		ax2.grid(True, alpha=0.3)
+
+		# ax3 = fig.add_subplot(133)
+		# ax3.set_ylabel('Integrated intensity')
+		# ax3.set_xlabel('Propagation distance [$\lambda$]')
+		# ax3.grid(True, alpha=0.3)
+		# ax3.tick_params(axis ='both', direction='in') 
+		# ax3.plot(Isum)
+		plt.tight_layout()
+		# plt.savefig('fresnel_cross.pdf', bbox_inches='tight')
+		# plt.savefig('fresnel_cross.png', bbox_inches='tight')
+		plt.show()
+
+	# E_xyz,_,_ = fou.planeWavePropagationFFTW(A, dz=4*l, dpix=dpix, oversample=1, show=True, save=False, outputLimits=(-0.4*1e-3,0.4*1e-3))#, scaleInputArray=(n,m))
+	# Exyz,_,_ = fou.FresnelPropagationFFTW(A, dz=4*l, dpix=dpix, oversample=1, show=True, save=False)#, scaleInputArray=(n,m))
+	end = time.time()
+	print((end-start)*1000)
+	# set_trace()
+
+

test cases (standalone)/fourierOptics_04_simple_lens_system.py

@@ -0,0 +1,720 @@
+#!/usr/bin/env Python3
+import numpy as np
+from pdb import set_trace
+import matplotlib.pyplot as plt
+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 time
+
+
+class FourierOptics:
+	def __init__(self, l):
+		self.l = l # Wavelength
+
+
+	def planeWavePropagationFFTW(self, A, dpix=None, dz=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.abs(_E_xy0)**2).sum()	# Factor to scale the intensity after the FFTs
+		_IntensityNormFactor = np.sum(np.abs(_E_xy0)**2)
+		_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)
+		_kz = self.phaseFactorExponent(_fx, _fy, self.l) 	# origin at start of array
+		_phaseFactor = np.exp(-1j * _kz * dz) 				# origin at start of array
+		
+		# Calculating the spatial frequencies in k-space
+		_kx 		= np.fft.fftshift( _fx * 2 * np.pi )
+		_ky 		= np.fft.fftshift( _fy * 2 * np.pi )
+
+		_E_xyz 		= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_A_in 		= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_A_in[:] 	= _Az0 * _phaseFactor
+
+		_FFTW_Obj2 	= pyfftw.FFTW( _A_in, _E_xyz, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
+		_FFTW_Obj2()
+		_E_xyz 		= np.fft.fftshift(_E_xyz) 				# use fftshift to move the origin from edge to center
+		_E_xyz /= np.sqrt(_IntensityNormFactor) 			# Normalization back to initial energy.
+
+		# Coordinates of E_xy0
+		_y, _x = np.mgrid[-_n / 2 * dpix : _n / 2 * dpix : 1j*_n, -_m / 2 * dpix : _m / 2 * dpix : 1j * _m]
+		if outputLimits:
+			# 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.
+			_x				= np.linspace(-_n/2*dpix, _n/2*dpix, _n)
+			_y				= np.linspace(-_m/2*dpix, _m/2*dpix, _m)
+			_f_ip_real		= interpolate.interp2d(_x, _y, _E_xyz.real, kind='linear')
+			_f_ip_imag 		= interpolate.interp2d(_x, _y, _E_xyz.imag, kind='linear')
+			_xi				= np.linspace(outputLimits[0],outputLimits[1],_n)
+			_yi 			= np.linspace(outputLimits[0],outputLimits[1],_m)
+			_E_out_ip_real 	= _f_ip_real(_xi, _yi)
+			_E_out_ip_imag	= _f_ip_imag(_xi, _yi)
+			_E_xyz 			= _E_out_ip_real + _E_out_ip_imag*1j
+			_y, _x 			= np.mgrid[outputLimits[0]:outputLimits[1]:1j*_n, outputLimits[0]:outputLimits[1]:1j * _m]
+			
+
+		_E_xyz_extent = [np.min(_x), np.max(_x), np.min(_y), np.max(_x)]
+
+
+
+		if show or save:
+			fig = plt.figure(figsize=(8,3))
+			ax = fig.add_subplot(121)
+			im = ax.imshow(np.abs(_E_xyz), cmap='hot', extent=_E_xyz_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_xyz), cmap='hot', extent=_E_xyz_extent)
+			ax2.set_xlabel('x [mm]')
+			ax2.set_ylabel('y [mm]')
+			fig.colorbar(im2, ax=ax2)
+			plt.tight_layout()
+		if save:
+			plt.savefig(save, bbox_inches='tight')
+		if show:
+			plt.show()
+		elif not show:
+			plt.close()
+
+		return np.array([_E_xyz, _x, _y])
+
+
+	def FresnelPropagationFFTW(self, A, xin, yin, dz=None, oversample=None, show=False, save=False, scaleInputArray=False, outputCoordinates='auto'):
+		'''
+			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]
+											If 'auto': the output coordinates will be determined 
+											automatically.
+
+			Output:
+			E_out 				- 	[array] (Complex) Output field distribution
+			x, y 				-	[array] Output coordinates 
+		'''
+		assert np.shape(xin) == np.shape(A), 'Input coordinates (xin) do not match shape of input field.'
+		assert np.shape(yin) == np.shape(A), 'Input coordinates (yin) do not match shape of input field.'
+		A = np.conj(A)
+
+		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='linear')
+			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)
+		_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
+		
+		# Calculate the pixel size from the xi, eta plane
+		_dpix = np.mean(np.diff(_xi))
+		if _dpix < 1e-99:
+			_dpix = np.mean(np.diff(_xi.T))
+
+		# Gaussian test: use 8*w0 radius for FFT:
+		if outputCoordinates == 'auto':
+			_w0_fit, _ 		= leastsq(self.minimize, self.l, args=(0, _xi, _eta, A))
+			_gauss 			= GaussianBeam(_w0_fit, self.l)
+			_halfwidth 		= 8 * _gauss.beamRadius(dz)[0]
+			_xip			= np.linspace(-_halfwidth, _halfwidth, _n) 						# Desired x,y - plane
+			_yip			= np.linspace(-_halfwidth, _halfwidth, _m) 						# Desired x,y - plane
+		else:
+			_xip			= np.linspace(outputCoordinates[0], outputCoordinates[1], _n) 	# Desired x,y - plane
+			_yip			= np.linspace(outputCoordinates[2], outputCoordinates[3], _m) 	# Desired x,y - plane
+
+
+		# The quadratic phase exponential of the input wavefront
+		_phaseComponent = np.exp(1j * _k / (2 * dz) * (_xi**2 + _eta**2))
+
+		# Complete input
+		_U_in 			= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_U_in[:]		= A * _phaseComponent
+		
+		# Output of the FFT evaluated at frequency fx = x / (lambda * z)
+		_U_out 			= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_FFTW_Obj		= pyfftw.FFTW( np.fft.ifftshift(_U_in), _U_out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
+		_FFTW_Obj()
+		_U_out 			= np.fft.fftshift(_U_out)
+
+		# Evaluate frequencies:
+		_fx 			= np.fft.fftshift(np.fft.fftfreq(np.shape(_xi)[0], d=_dpix))
+		_fy 			= np.fft.fftshift(np.fft.fftfreq(np.shape(_xi)[1], d=_dpix))
+		_x 				= _fx * self.l * dz
+		_y 				= _fy * self.l * dz
+		_x2, _y2 		= np.meshgrid(_x, _y)
+
+
+		# Final output and frequencies
+		_pre_amble 		= np.exp(1j * _k * dz) / (1j * self.l * dz) * np.exp(1j * _k / (2 * dz) * (_x2**2 + _y2**2))
+		_E_out 			= _pre_amble * _U_out
+
+		# Normalization of the output values
+		_IntensityNormFactor = np.sum(np.abs(A)**2)
+		_E_out /= (np.abs(np.exp(1j*_k*dz)/(1j*self.l*dz)) * np.sqrt(_IntensityNormFactor) * np.sqrt(_n*_m))
+
+
+		# Performing interpolation to get to the required final plane, done for real and imaginary part.
+		_f_ip_real		= interpolate.interp2d(_x.real, _y.real, _E_out.real, kind='linear')
+		_f_ip_imag 		= interpolate.interp2d(_x.real, _y.real, _E_out.imag, kind='linear')
+		_E_out_ip_real 	= _f_ip_real(_xip, _yip)
+		_E_out_ip_imag	= _f_ip_imag(_xip, _yip)
+		_E_out_ip 		= _E_out_ip_real + _E_out_ip_imag*1j
+		_E_out_ip 		/= np.sqrt(((_x2.max()/_xip.max())*(_y2.max()/_yip.max())))
+		_E_out_ip = np.conj(_E_out_ip)
+
+		if show:
+			# Plotting the results
+			_fig = plt.figure(figsize=(8,3))
+			_ax = _fig.add_subplot(121)
+			_im = _ax.imshow(np.abs(_E_out_ip), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()])
+			_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='Amplitude')
+			_ax2 = _fig.add_subplot(122)
+			_im2 = _ax2.imshow(np.angle(_E_out_ip), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()])
+			_ax2.set_xlim([_xip.min(), _xip.max()])
+			_ax2.set_ylim([_yip.min(), _yip.max()])
+			_ax2.set_xlabel('x [mm]')
+			_ax2.set_ylabel('y [mm]')
+			_fig.colorbar(_im2, ax=_ax2, label='Phase')
+			plt.tight_layout()
+			if save:
+				plt.savefig(save, bbox_inches='tight', transparent=True)
+			plt.show()
+
+		_x, _y = np.meshgrid(_xip, _yip)
+		return np.array([_E_out_ip, _x, _y])
+
+
+	def freqArray(self, A, dpix):
+		n, m = np.shape(A)
+		
+		fy, fx = np.mgrid[0:m, 0:n]
+
+		fx = fx.flatten()
+		fy = fy.flatten()
+		
+		for i in range(len(fx)):
+			fx[i] = (fx[i] - n*bool(fx[i]>(n-1)/2) ) / (n*dpix)
+			fy[i] = (fy[i] - m*bool(fy[i]>(m-1)/2) ) / (m*dpix)
+		
+		fx = fx.reshape(n,m)
+		fy = fy.reshape(n,m)
+		
+		return fx, fy
+				
+
+	def phaseFactorExponent(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 minimize(self, a, dz, x, y, og_data):
+		'''
+			Residual function used for fitting the gaussian beam to the input distribution.
+			Input:
+			a 		- 	[float] Guess for w0
+			dz  	- 	[float] Propagation distance, should be known
+			x,y 	- 	[array] meshgrid of the x and y coordinates.
+			og_data - 	[array] input distribution to be fitted.
+		'''
+		return np.abs((self.gaussianBeamFit(a,dz,x,y) - og_data).flatten()).real
+
+
+	def gaussianBeamFit(self, a, dz, x, y):
+		'''
+			Function used to fit the Gaussian beam to an input plane. Very similar
+			to the GaussianBeam class below, but modified to allow a best fit.
+			Input:
+			a 	- 	[float] Guess for w0
+			dz  - 	[float] Propagation distance, should be known
+			x,y - 	[array] meshgrid of the x and y coordinates.
+		'''
+		_w0 = a
+		_z 	= dz
+		_r 	  	= np.sqrt(x**2 + y**2)
+		if _z != 0:
+			_z_c 	= np.pi * _w0**2 / self.l
+			_w 		= _w0 * np.sqrt(1 + (_z/_z_c)**2)
+			_R 		= _z + _z_c **2 / _z
+			_phi0 	= np.arctan(_z/_z_c)
+			_k 		= 2 * np.pi / self.l
+			_results = np.sqrt(2 / (np.pi*_w**2)) * np.exp(- _r**2 / _w**2 - 1j * _k * _z - 1j * np.pi * _r**2 / (self.l * _R) + 1j*_phi0)
+		elif _z == 0:
+			_results = np.sqrt(2 / (np.pi*_w0**2)) * np.exp( - _r**2 / _w0**2 )
+		return _results
+
+
+
+
+
+
+class profiles:
+	def __init__(self):
+		pass
+
+	def topHat(self, r, n, m, dpix, oversample=1):
+	    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 <= (r)).astype(int)
+	    return _Anew
+
+	def obscuredTopHat(self, r1, r2, n, m, dpix, oversample=1):
+		'''
+			Creates a grid of n x m points, following a top hat function
+			where all points inside a radius r is 1 and 0 otherwise.
+		'''
+		assert r1 > r2, "r2 < r1, change input"
+		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))
+		return _Anew.astype(int)
+
+
+	def spiderTopHat(self, r1, r2, n, m, dpix, oversample=1):
+		'''
+			Creates a grid of n x m points, following a top hat function
+			where all points inside a radius r is 1 and 0 otherwise.
+		'''
+		assert r1 > r2, "r2 < r1, change input"
+		j,i 				= np.mgrid[-n/2*dpix : n/2*dpix : 1j*n*oversample, -m/2*dpix : m/2*dpix : 1j*m*oversample]
+		n *= oversample
+		m *= oversample
+		dist_from_center 	= np.sqrt(i**2 + j**2)
+		_Anew 				= (dist_from_center <= (r1))  & (dist_from_center >= (r2))
+		_Aspider 			= (np.abs(i) < (0.05*r1)) | (np.abs(j) < (0.05*r1))
+		_Anew 				= _Anew & ~_Aspider
+		return _Anew
+
+
+
+class GaussianBeam:
+	def __init__(self, w0, l):
+		# Confocal distance
+		self.z_c 	 = np.pi * w0**2 / l
+
+		# Far-field divergence angle
+		self.theta_0 = l / (np.pi * w0)
+		self.l 		 = l
+		self.w0 	 = w0
+		self.fwhm 	 = 1.18 * self.theta_0
+
+	def fundamentalMode(self, z, n, m, rmax):
+		y, x = np.mgrid[-rmax:rmax:n*1j,-rmax:rmax:n*1j]
+		r 	 = np.sqrt(x**2 + y**2)
+		if z != 0:
+			_w 		= self.beamRadius(z)
+			_R 		= self.curvature(z)
+			_phi0 	= self.phaseShift(z)
+			_k 		= 2 * np.pi / self.l
+			return np.sqrt(2 / (np.pi*_w**2)) * np.exp(- r**2 / _w**2 - 1j * _k * z - 1j * np.pi * r**2 / (self.l * _R) + 1j*_phi0)
+		elif z == 0:
+			return np.sqrt(2 / (np.pi*self.w0**2)) * np.exp( - r**2 / self.w0**2 ) 
+
+	def curvature(self, z):
+		return z + self.z_c **2 / z
+
+	def beamRadius(self, z):
+		return self.w0 * np.sqrt(1 + (z/self.z_c)**2)
+
+	def phaseShift(self, z):
+		return np.arctan(z/self.z_c)
+
+	def relativePower(self, r, z):
+		return np.exp(-2 * (r / self.beamRadius(z))**2)
+
+
+
+class lensObject:
+	'''
+		The lens object describes any spherical lens and find the phase delay
+		for all light traveling through it.
+		The refractive index, of course, depends on the wavelength.
+	'''
+	def __init__(self, glass, R1, R2, thickness, CA, f=None, name=None):
+		self.glass 		= glass
+		assert glassCatalog.catalog.get(self.glass), 'Glass type not in catalog'
+		self.R1 		= R1 				# mm
+		self.R2 		= R2 				# mm
+		self.thickness 	= thickness			# mm
+		self.CA 		= CA 				# mm
+		self.dispersionCoefficients = glassCatalog.catalog.get(self.glass)
+		self.referenceWavelength = 1.35 	# micron
+		self.type 		= 'Spherical lens'
+		if not f:
+			self.f 		= self.findFocalLength()
+		else:
+			self.f 		= f
+			self.type 	= 'Thin paraxial lens'
+		if name:
+			assert type(name)==str, 'No string for a name.'
+			self.name = name
+		else:
+			self.name = str(self)
+
+	def findRefractiveIndex(self, _l):
+		_a = self.dispersionCoefficients
+		_n = np.sqrt( (_a[0]* _l**2)/(_l**2 - _a[1]) + (_a[2] * _l**2) / (_l**2 - _a[3]) + (_a[4] * _l**2) / (_l**2 - _a[5]) + 1)
+		return _n
+
+	def findFocalLength(self):
+		_n 		= self.findRefractiveIndex(self.referenceWavelength)
+		_f_recip = (_n - 1) * (1 / self.R1 - 1 / self.R2)
+		return 1/_f_recip
+
+	def phase_delay(self, x, y, l, paraxial=False):
+		if self.type == 'Thin paraxial lens':
+			assert self.type=='Thin paraxial lens' and paraxial, 'Can not calculate non-paraxial phase delay on paraxial lens!'
+		
+		_k = 2 * np.pi / l
+		_n = self.findRefractiveIndex(self.referenceWavelength)
+
+		if not paraxial:
+			_A1 = self.R1 * (1 - np.sqrt(1 - (x**2 + y**2) / self.R1**2 ) )
+			_A2 = self.R2 * (1 - np.sqrt(1 - (x**2 + y**2) / self.R2**2 ) )
+			_dxy = self.thickness - _A1 + _A2
+			_t_l = np.exp(-1j * _k * self.thickness) * np.exp(-1j * _k * (_n - 1) * _dxy)
+		if paraxial:
+			_t_l = np.exp(-1j * _k * self.thickness) * np.exp(1j * _k / (2 * self.f) * (x**2 + y**2))
+
+		# 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)).astype(int)
+		_t_l *= _lensMask
+		return _t_l
+
+
+
+class mirrorObject:
+	'''
+		Describes any mirror. No wavelength dependence exist. 
+		Coatings are not considered.
+	'''
+	def __init__(self, dx, dy, name=None):
+		self.dx = dx
+		self.dy = dy
+		if name:
+			assert type(name)==str, 'No string for a name.'
+			self.name = name
+		else:
+			self.name = str(self)
+
+
+class parabolicMirror(mirrorObject):
+	def __init__(self, f, CA, dx=None, dy=None, name=None):
+		mirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.f = f
+		self.CA = CA
+		self.type = 'Parabolic mirror'
+
+	def phase_delay(self, x, y, l):
+		_k 		= 2 * np.pi / l
+		_dz 	= (x**2 + y**2) / (4 * self.f)
+		_t_m 	= np.exp(-1j * _k * 2 * _dz)
+		return _t_m
+
+
+
+class sphericalMirror(mirrorObject):
+	def __init__(self, R, CA, dx=None, dy=None, name=None):
+		mirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.R = R
+		self.CA = CA
+		self.type = 'Spherical mirror'
+
+	def phase_delay(self, x, y, l):
+		_k = 2 * np.pi / l
+		_dz = self.R * (1 - np.sqrt(1 - (x**2 + y**2)/self.R**2) )
+		_t_m = np.exp(-1j * _k * 2 * _dz)
+		return _t_m
+
+
+
+class freeformMirror(mirrorObject):
+	def __init__(self, x, y, z, dx=None, dy=None, name=None):
+		mirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.x = x
+		self.y = y
+		self.z = z
+		self.type = 'User def. mirror'
+
+	def phase_delay(self, l):
+		_k = 2 * np.pi / l
+		_t_m = np.exp(-1j * _k * 2 * self.z)
+		return _t_m
+
+class opticalSurface:
+	def __init__(self, mask=None, name=None):
+		self.type = 'Dummy surface'
+		if mask:
+			self.x, self.y, self.mask = mask
+			self.type = 'Mask'
+		if name:
+			assert type(name)==str, 'No string for a name.'
+			self.name = name
+
+		pass
+
+class OpticalTrain:
+	'''
+		The main class where an optical system is build up sequentially.
+	'''
+	def __init__(self):
+		self.opticalTrain 		= []
+		self.sequentialDistance = []
+
+	def add_object(self, obj, dist):
+		'''
+			Adding a lens or mirror object to the opticalTrain
+			obj 	The lens or mirror object
+			dist 	The distance from the the previous object.
+		'''
+		assert type(dist) == float or type(dist) == int, 'Distance is not a float or integer.'
+		self.opticalTrain.append(obj)
+		self.sequentialDistance.append(float(dist))
+
+
+	def print_lens_Data(self):
+		'''
+			Print the data of the current opticalTrain.
+		'''
+		_tableData = []
+		for _lens in self.opticalTrain:
+			_i = self.opticalTrain.index(_lens)
+			_distances = self.sequentialDistance
+			_distances.append('-')
+			_tableData.append([_i, _lens.name, _lens.type, _distances[_i+1] ])
+		print(tabulate(_tableData, headers=['#', 'Name', 'Type', 'Distance to next surface']))
+
+	def fresnel_number(self, a, z, l):
+		return a**2 / (z * l)
+
+	def propagate_system(self, U_in, x, y, l, output=None, paraxial=True):
+		# We need to have an input distribution
+		# We need to determine what propagation method is required
+		# Call the objects in the self.opticalTrain sequentially
+		# Call the subsequent phaseDelay
+		if output == 'complete':
+			_wf_list 	= []
+			_x_list 	= []
+			_y_list 	= []
+			_wf_list.append(_wf)
+			_x_list.append(x)
+			_y_list.append(y)
+		
+		_fou = FourierOptics(l=l)
+
+		_wf = U_in
+
+		oC = [x.min(), x.max(), y.min(), y.max()]
+		show = False
+		_wf, x, y = _fou.FresnelPropagationFFTW(_wf, x, y, dz=self.sequentialDistance[0], oversample=None, show=show, save=False, outputCoordinates=oC)
+		if output == 'complete':
+			_wf_list.append(_wf)
+			_x_list.append(x)
+			_y_list.append(y)
+		
+		
+		for _lens in self.opticalTrain[:-1]: # Except for image plane
+
+			title = f"{self.opticalTrain[:-1].index(_lens)+1:02d}"
+			print (' Calculating lens {0} of {1}'.format(title, f"{len(self.opticalTrain[:-1]):02d}"), end="\r")
+
+			_i 	= self.opticalTrain.index(_lens)
+			_dz = self.sequentialDistance[_i+1]
+			_wf *= _lens.phase_delay(x, y, l, paraxial=paraxial)
+
+			# Perform Fresnel number test:
+			_F = self.fresnel_number(_lens.CA, _dz, l)
+			_F = 1
+
+			
+			if _F < 5:
+				# Fresnel propagation
+				if _lens == self.opticalTrain[-2]:
+					oC = [-2, 2, -2, 2]
+					show=True
+				_wf, x, y = _fou.FresnelPropagationFFTW(_wf, x, y, dz=self.sequentialDistance[_i+1], oversample=None, show=show, save=False, outputCoordinates=oC)
+			elif _F >= 5:
+				print('Using Plane Wave!')
+				# Angular spectrum (p2p) propagation
+				# _wf, x, y = self.plane_wave_propagation()
+
+			pass
+			if output == 'complete':
+				_wf_list.append(_wf)
+				_x_list.append(x)
+				_y_list.append(y)
+		
+		if output == 'Complete':
+			return np.array([_wf_list, _x_list, _y_list])
+
+		else:
+			return np.array([_wf, x, y])
+
+
+		
+
+
+
+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/20
+	n, m = 4001, 4001
+	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)
+	A = profiles().topHat(25, 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))
+	
+	# Start an optical train object. Add lenses and propagate. Currently not in use.
+	LDE = OpticalTrain()
+	# Create a lens with D=50 mm, f=200 mm
+	L1 = lensObject('N-BK7', None, None, 0, 25, f=200, name='L1')
+	L2 = lensObject('N-BK7', None, None, 0, 12.5, f=-100, name='L2')
+	L3 = lensObject('N-BK7', None, None, 0, 25, f=200, name='L3')
+	IMG = opticalSurface(name='Image plane')
+
+	L0 = lensObject('N-BK7', None, None, 0, 25, f=400, name='Input intiater lens')
+	t_l = L0.phase_delay(x,y,l, paraxial=True)
+	U_in = A * t_l
+	set_trace()
+	# U_in = Fresnel propagation to L0.f
+	
+	oC = [x.min(), x.max(), y.min(), y.max()]
+	Exyz, x1, y1 = fou.FresnelPropagationFFTW(U_in, x, y, dz=L0.f, oversample=None, show=False, save=False, outputCoordinates=oC)
+	U_in = Exyz
+		
+
+	LDE.add_object(L1, 400)
+	LDE.add_object(L2, 200)
+	LDE.add_object(L3, 200)
+	LDE.add_object(IMG, 400)
+
+	LDE.print_lens_Data()
+	LDE.propagate_system(U_in, x, y, l)
+	# fig = plt.figure()
+	# ax1 = fig.add_subplot(121)
+	# ax2 = fig.add_subplot(122)
+	# im1 = ax1.imshow(np.abs(wf), cmap='hot', extent=[x.min(), x.max(), y.min(), y.max()])
+	# im2 = ax2.imshow(np.angle(wf), cmap='hot', extent=[x.min(), x.max(), y.min(), y.max()])
+	# plt.colorbar(im1, ax=ax1)
+	# plt.colorbar(im2, ax=ax2)
+	# plt.show()
+	set_trace()
+	
+
+
+
+
+	# Closing down the script and stopping the timer.
+	end = time.time()
+	print('The program took {:.3f} seconds!'.format((end-start)))
+	
+
+

test cases (standalone)/fourierOptics_25_single_lens_system.py

@@ -0,0 +1,729 @@
+#!/usr/bin/env Python3
+import numpy as np
+from pdb import set_trace
+import matplotlib.pyplot as plt
+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 time
+
+
+class FourierOptics:
+	def __init__(self, l):
+		self.l = l # Wavelength
+
+
+	def planeWavePropagationFFTW(self, A, dpix=None, dz=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.abs(_E_xy0)**2).sum()	# Factor to scale the intensity after the FFTs
+		_IntensityNormFactor = np.sum(np.abs(_E_xy0)**2)
+		_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)
+		_kz = self.phaseFactorExponent(_fx, _fy, self.l) 	# origin at start of array
+		_phaseFactor = np.exp(-1j * _kz * dz) 				# origin at start of array
+		
+		# Calculating the spatial frequencies in k-space
+		_kx 		= np.fft.fftshift( _fx * 2 * np.pi )
+		_ky 		= np.fft.fftshift( _fy * 2 * np.pi )
+
+		_E_xyz 		= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_A_in 		= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_A_in[:] 	= _Az0 * _phaseFactor
+
+		_FFTW_Obj2 	= pyfftw.FFTW( _A_in, _E_xyz, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
+		_FFTW_Obj2()
+		_E_xyz 		= np.fft.fftshift(_E_xyz) 				# use fftshift to move the origin from edge to center
+		_E_xyz /= np.sqrt(_IntensityNormFactor) 			# Normalization back to initial energy.
+
+		# Coordinates of E_xy0
+		_y, _x = np.mgrid[-_n / 2 * dpix : _n / 2 * dpix : 1j*_n, -_m / 2 * dpix : _m / 2 * dpix : 1j * _m]
+		if outputLimits:
+			# 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.
+			_x				= np.linspace(-_n/2*dpix, _n/2*dpix, _n)
+			_y				= np.linspace(-_m/2*dpix, _m/2*dpix, _m)
+			_f_ip_real		= interpolate.interp2d(_x, _y, _E_xyz.real, kind='linear')
+			_f_ip_imag 		= interpolate.interp2d(_x, _y, _E_xyz.imag, kind='linear')
+			_xi				= np.linspace(outputLimits[0],outputLimits[1],_n)
+			_yi 			= np.linspace(outputLimits[0],outputLimits[1],_m)
+			_E_out_ip_real 	= _f_ip_real(_xi, _yi)
+			_E_out_ip_imag	= _f_ip_imag(_xi, _yi)
+			_E_xyz 			= _E_out_ip_real + _E_out_ip_imag*1j
+			_y, _x 			= np.mgrid[outputLimits[0]:outputLimits[1]:1j*_n, outputLimits[0]:outputLimits[1]:1j * _m]
+			
+
+		_E_xyz_extent = [np.min(_x), np.max(_x), np.min(_y), np.max(_x)]
+
+
+
+		if show or save:
+			fig = plt.figure(figsize=(8,3))
+			ax = fig.add_subplot(121)
+			im = ax.imshow(np.abs(_E_xyz), cmap='hot', extent=_E_xyz_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_xyz), cmap='hot', extent=_E_xyz_extent)
+			ax2.set_xlabel('x [mm]')
+			ax2.set_ylabel('y [mm]')
+			fig.colorbar(im2, ax=ax2)
+			plt.tight_layout()
+		if save:
+			plt.savefig(save, bbox_inches='tight')
+		if show:
+			plt.show()
+		elif not show:
+			plt.close()
+
+		return np.array([_E_xyz, _x, _y])
+
+
+	def FresnelPropagationFFTW(self, A, xin, yin, dz=None, oversample=None, show=False, save=False, scaleInputArray=False, outputCoordinates='auto'):
+		'''
+			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]
+											If 'auto': the output coordinates will be determined 
+											automatically.
+
+			Output:
+			E_out 				- 	[array] (Complex) Output field distribution
+			x, y 				-	[array] Output coordinates 
+		'''
+		assert np.shape(xin) == np.shape(A), 'Input coordinates (xin) do not match shape of input field.'
+		assert np.shape(yin) == np.shape(A), 'Input coordinates (yin) do not match shape of input field.'
+		A = np.conj(A)
+
+		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='linear')
+			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)
+		_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
+		
+		# Calculate the pixel size from the xi, eta plane
+		_dpix = np.mean(np.diff(_xi))
+		if _dpix < 1e-99:
+			_dpix = np.mean(np.diff(_xi.T))
+
+		# Gaussian test: use 8*w0 radius for FFT:
+		if outputCoordinates == 'auto':
+			_w0_fit, _ 		= leastsq(self.minimize, self.l, args=(0, _xi, _eta, A))
+			_gauss 			= GaussianBeam(_w0_fit, self.l)
+			_halfwidth 		= 8 * _gauss.beamRadius(dz)[0]
+			_xip			= np.linspace(-_halfwidth, _halfwidth, _n) 						# Desired x,y - plane
+			_yip			= np.linspace(-_halfwidth, _halfwidth, _m) 						# Desired x,y - plane
+		else:
+			_xip			= np.linspace(outputCoordinates[0], outputCoordinates[1], _n) 	# Desired x,y - plane
+			_yip			= np.linspace(outputCoordinates[2], outputCoordinates[3], _m) 	# Desired x,y - plane
+
+
+		# The quadratic phase exponential of the input wavefront
+		_phaseComponent = np.exp(1j * _k / (2 * dz) * (_xi**2 + _eta**2))
+
+		# Complete input
+		_U_in 			= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_U_in[:]		= A * _phaseComponent
+		
+		# Output of the FFT evaluated at frequency fx = x / (lambda * z)
+		_U_out 			= pyfftw.empty_aligned((_n,_m), dtype='complex128')
+		_FFTW_Obj		= pyfftw.FFTW( np.fft.ifftshift(_U_in), _U_out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
+		_FFTW_Obj()
+		_U_out 			= np.fft.fftshift(_U_out)
+
+		# Evaluate frequencies:
+		_fx 			= np.fft.fftshift(np.fft.fftfreq(np.shape(_xi)[0], d=_dpix))
+		_fy 			= np.fft.fftshift(np.fft.fftfreq(np.shape(_xi)[1], d=_dpix))
+		_x 				= _fx * self.l * dz
+		_y 				= _fy * self.l * dz
+		_x2, _y2 		= np.meshgrid(_x, _y)
+
+
+		# Final output and frequencies
+		_pre_amble 		= np.exp(1j * _k * dz) / (1j * self.l * dz) * np.exp(1j * _k / (2 * dz) * (_x2**2 + _y2**2))
+		_E_out 			= _pre_amble * _U_out
+
+		# Normalization of the output values
+		_IntensityNormFactor = np.sum(np.abs(A)**2)
+		_E_out /= (np.abs(np.exp(1j*_k*dz)/(1j*self.l*dz)) * np.sqrt(_IntensityNormFactor) * np.sqrt(_n*_m))
+
+		# Performing interpolation to get to the required final plane, done for real and imaginary part.
+		_f_ip_real		= interpolate.interp2d(_x, _y, _E_out.real, kind='linear')
+		_f_ip_imag 		= interpolate.interp2d(_x, _y, _E_out.imag, kind='linear')
+		_E_out_ip_real 	= _f_ip_real(_xip, _yip)
+		_E_out_ip_imag	= _f_ip_imag(_xip, _yip)
+		_E_out_ip 		= _E_out_ip_real + _E_out_ip_imag*1j
+		_E_out_ip 		/= np.sqrt(((_x2.max()/_xip.max())*(_y2.max()/_yip.max())))
+		_E_out_ip = np.conj(_E_out_ip)
+
+		if show:
+			# Plotting the results
+			_fig = plt.figure(figsize=(8,3))
+			_ax = _fig.add_subplot(121)
+			_im = _ax.imshow(np.abs(_E_out_ip), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()])
+			_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='Amplitude')
+			_ax2 = _fig.add_subplot(122)
+			_im2 = _ax2.imshow(np.angle(_E_out_ip), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()])
+			_ax2.set_xlim([_xip.min(), _xip.max()])
+			_ax2.set_ylim([_yip.min(), _yip.max()])
+			_ax2.set_xlabel('x [mm]')
+			_ax2.set_ylabel('y [mm]')
+			_fig.colorbar(_im2, ax=_ax2, label='Phase')
+			plt.tight_layout()
+			if save:
+				plt.savefig(save, bbox_inches='tight', transparent=True)
+			plt.show()
+
+		_x, _y = np.meshgrid(_xip, _yip)
+		return np.array([_E_out_ip, _x, _y])
+
+
+	def freqArray(self, A, dpix):
+		n, m = np.shape(A)
+		
+		fy, fx = np.mgrid[0:m, 0:n]
+
+		fx = fx.flatten()
+		fy = fy.flatten()
+		
+		for i in range(len(fx)):
+			fx[i] = (fx[i] - n*bool(fx[i]>(n-1)/2) ) / (n*dpix)
+			fy[i] = (fy[i] - m*bool(fy[i]>(m-1)/2) ) / (m*dpix)
+		
+		fx = fx.reshape(n,m)
+		fy = fy.reshape(n,m)
+		
+		return fx, fy
+				
+
+	def phaseFactorExponent(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 minimize(self, a, dz, x, y, og_data):
+		'''
+			Residual function used for fitting the gaussian beam to the input distribution.
+			Input:
+			a 		- 	[float] Guess for w0
+			dz  	- 	[float] Propagation distance, should be known
+			x,y 	- 	[array] meshgrid of the x and y coordinates.
+			og_data - 	[array] input distribution to be fitted.
+		'''
+		return np.abs((self.gaussianBeamFit(a,dz,x,y) - og_data).flatten()).real
+
+
+	def gaussianBeamFit(self, a, dz, x, y):
+		'''
+			Function used to fit the Gaussian beam to an input plane. Very similar
+			to the GaussianBeam class below, but modified to allow a best fit.
+			Input:
+			a 	- 	[float] Guess for w0
+			dz  - 	[float] Propagation distance, should be known
+			x,y - 	[array] meshgrid of the x and y coordinates.
+		'''
+		_w0 = a
+		_z 	= dz
+		_r 	  	= np.sqrt(x**2 + y**2)
+		if _z != 0:
+			_z_c 	= np.pi * _w0**2 / self.l
+			_w 		= _w0 * np.sqrt(1 + (_z/_z_c)**2)
+			_R 		= _z + _z_c **2 / _z
+			_phi0 	= np.arctan(_z/_z_c)
+			_k 		= 2 * np.pi / self.l
+			_results = np.sqrt(2 / (np.pi*_w**2)) * np.exp(- _r**2 / _w**2 - 1j * _k * _z - 1j * np.pi * _r**2 / (self.l * _R) + 1j*_phi0)
+		elif _z == 0:
+			_results = np.sqrt(2 / (np.pi*_w0**2)) * np.exp( - _r**2 / _w0**2 )
+		return _results
+
+
+
+
+
+
+class profiles:
+	def __init__(self):
+		pass
+
+	def topHat(self, r, n, m, dpix, oversample=1):
+	    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 <= (r)).astype(int)
+	    return _Anew
+
+	def obscuredTopHat(self, r1, r2, n, m, dpix, oversample=1):
+		'''
+			Creates a grid of n x m points, following a top hat function
+			where all points inside a radius r is 1 and 0 otherwise.
+		'''
+		assert r1 > r2, "r2 < r1, change input"
+		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))
+		return _Anew.astype(int)
+
+
+	def spiderTopHat(self, r1, r2, n, m, dpix, oversample=1):
+		'''
+			Creates a grid of n x m points, following a top hat function
+			where all points inside a radius r is 1 and 0 otherwise.
+		'''
+		assert r1 > r2, "r2 < r1, change input"
+		j,i 				= np.mgrid[-n/2*dpix : n/2*dpix : 1j*n*oversample, -m/2*dpix : m/2*dpix : 1j*m*oversample]
+		n *= oversample
+		m *= oversample
+		dist_from_center 	= np.sqrt(i**2 + j**2)
+		_Anew 				= (dist_from_center <= (r1))  & (dist_from_center >= (r2))
+		_Aspider 			= (np.abs(i) < (0.05*r1)) | (np.abs(j) < (0.05*r1))
+		_Anew 				= _Anew & ~_Aspider
+		return _Anew
+
+
+
+class GaussianBeam:
+	def __init__(self, w0, l):
+		# Confocal distance
+		self.z_c 	 = np.pi * w0**2 / l
+
+		# Far-field divergence angle
+		self.theta_0 = l / (np.pi * w0)
+		self.l 		 = l
+		self.w0 	 = w0
+		self.fwhm 	 = 1.18 * self.theta_0
+
+	def fundamentalMode(self, z, n, m, rmax):
+		y, x = np.mgrid[-rmax:rmax:n*1j,-rmax:rmax:n*1j]
+		r 	 = np.sqrt(x**2 + y**2)
+		if z != 0:
+			_w 		= self.beamRadius(z)
+			_R 		= self.curvature(z)
+			_phi0 	= self.phaseShift(z)
+			_k 		= 2 * np.pi / self.l
+			return np.sqrt(2 / (np.pi*_w**2)) * np.exp(- r**2 / _w**2 - 1j * _k * z - 1j * np.pi * r**2 / (self.l * _R) + 1j*_phi0)
+		elif z == 0:
+			return np.sqrt(2 / (np.pi*self.w0**2)) * np.exp( - r**2 / self.w0**2 ) 
+
+	def curvature(self, z):
+		return z + self.z_c **2 / z
+
+	def beamRadius(self, z):
+		return self.w0 * np.sqrt(1 + (z/self.z_c)**2)
+
+	def phaseShift(self, z):
+		return np.arctan(z/self.z_c)
+
+	def relativePower(self, r, z):
+		return np.exp(-2 * (r / self.beamRadius(z))**2)
+
+
+
+class lensObject:
+	'''
+		The lens object describes any spherical lens and find the phase delay
+		for all light traveling through it.
+		The refractive index, of course, depends on the wavelength.
+	'''
+	def __init__(self, glass, R1, R2, thickness, CA, f=None, name=None):
+		self.glass 		= glass
+		assert glassCatalog.catalog.get(self.glass), 'Glass type not in catalog'
+		self.R1 		= R1 				# mm
+		self.R2 		= R2 				# mm
+		self.thickness 	= thickness			# mm
+		self.CA 		= CA 				# mm
+		self.dispersionCoefficients = glassCatalog.catalog.get(self.glass)
+		self.referenceWavelength = 1.35 	# micron
+		self.type 		= 'Spherical lens'
+		if not f:
+			self.f 		= self.findFocalLength()
+		else:
+			self.f 		= f
+		if name:
+			assert type(name)==str, 'No string for a name.'
+			self.name = name
+		else:
+			self.name = str(self)
+
+	def findRefractiveIndex(self, _l):
+		_a = self.dispersionCoefficients
+		_n = np.sqrt( (_a[0]* _l**2)/(_l**2 - _a[1]) + (_a[2] * _l**2) / (_l**2 - _a[3]) + (_a[4] * _l**2) / (_l**2 - _a[5]) + 1)
+		return _n
+
+	def findFocalLength(self):
+		_n 		= self.findRefractiveIndex(self.referenceWavelength)
+		_f_recip = (_n - 1) * (1 / self.R1 - 1 / self.R2)
+		return 1/_f_recip
+
+	def phaseDelayFunction(self, x, y, l, paraxial=False):
+		if not paraxial:
+			_A1 = self.R1 * (1 - np.sqrt(1 - (x**2 + y**2) / self.R1**2 ) )
+			_A2 = self.R2 * (1 - np.sqrt(1 - (x**2 + y**2) / self.R2**2 ) )
+			_dxy = self.thickness - _A1 + _A2
+		if paraxial:
+			_dxy = self.thickness - (x**2 + y**2)/2 * (1/self.R1 - 1/self.R2)
+		_k = 2 * np.pi / 30e-3
+		_n = self.findRefractiveIndex(self.referenceWavelength)
+		_t_l = np.exp(-1j * _k * self.thickness)* np.exp(-1j * _k * (_n - 1) * _dxy)
+
+		# 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)).astype(int)
+		_t_l *= _lensMask
+		return _t_l
+
+
+
+class mirrorObject:
+	'''
+		Describes any mirror. No wavelength dependence exist. 
+		Coatings are not considered.
+	'''
+	def __init__(self, dx, dy, name=None):
+		self.dx = dx
+		self.dy = dy
+		if name:
+			assert type(name)==str, 'No string for a name.'
+			self.name = name
+		else:
+			self.name = str(self)
+
+
+class parabolicMirror(mirrorObject):
+	def __init__(self, f, CA, dx=None, dy=None, name=None):
+		mirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.f = f
+		self.CA = CA
+		self.type = 'Parabolic mirror'
+
+	def phaseDelayFunction(self, x, y, l):
+		_k 		= 2 * np.pi / l
+		_dz 	= (x**2 + y**2) / (4 * self.f)
+		_t_m 	= np.exp(-1j * _k * 2 * _dz)
+		return _t_m
+
+
+
+class sphericalMirror(mirrorObject):
+	def __init__(self, R, CA, dx=None, dy=None, name=None):
+		mirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.R = R
+		self.CA = CA
+		self.type = 'Spherical mirror'
+
+	def phaseDelayFunction(self, x, y, l):
+		_k = 2 * np.pi / l
+		_dz = self.R * (1 - np.sqrt(1 - (x**2 + y**2)/self.R**2) )
+		_t_m = np.exp(-1j * _k * 2 * _dz)
+		return _t_m
+
+
+
+class freeformMirror(mirrorObject):
+	def __init__(self, x, y, z, dx=None, dy=None, name=None):
+		mirrorObject.__init__(self, dx=dx, dy=dy, name=name)
+		self.x = x
+		self.y = y
+		self.z = z
+		self.type = 'User def. mirror'
+
+	def phaseDelayFunction(self, l):
+		_k = 2 * np.pi / l
+		_t_m = np.exp(-1j * _k * 2 * self.z)
+		return _t_m
+
+
+
+
+class OpticalTrain:
+	'''
+		The main class where an optical system is build up sequentially.
+	'''
+	def __init__(self):
+		self.opticalTrain 		= []
+		self.sequentialDistance = []
+
+	def add_object(self, obj, dist):
+		'''
+			Adding a lens or mirror object to the opticalTrain
+			obj 	The lens or mirror object
+			dist 	The distance from the the previous object.
+		'''
+		assert type(dist) == float or type(dist) == int, 'Distance is not a float or integer.'
+		self.opticalTrain.append(obj)
+		self.sequentialDistance.append(float(dist))
+
+	def print_lens_Data(self):
+		'''
+			Print the data of the current opticalTrain.
+		'''
+		_tableData = []
+		for _lens in self.opticalTrain:
+			_i = self.opticalTrain.index(_lens)
+			_distances = self.sequentialDistance
+			_distances.append('-')
+			_tableData.append([_i, _lens.name, _lens.type, _distances[_i+1] ])
+		print(tabulate(_tableData, headers=['#', 'Name', 'Type', 'Distance to next surface']))
+
+
+	def propagate_system(self, U_in, x, y, l, output=None, paraxial=True):
+		# We need to have an input distribution
+		# We need to determine what propagation method is required
+		# Call the objects in the self.opticalTrain sequentially
+		# Call the subsequent phaseDelay
+		if output == 'complete':
+			_wf_list 	= []
+			_x_list 	= []
+			_y_list 	= []
+			_wf_list.append(_wf)
+			_x_list.append(x)
+			_y_list.append(y)
+		
+		_wf = U_in
+
+		for _lens in self.opticalTrain:
+			_i 	= self.opticalTrain.index(_lens)
+			_dz = self.sequentialDistance[_i]
+			_wf = _lens.phaseDelayFunction(x, y, l, paraxial=paraxial)
+
+			# Perform Fresnel number test:
+			_F = self.fresnel_number()
+
+			if _F < 5:
+				# Fresnel propagation
+				_wf, x, y = self.fresnel_propagation()
+			elif _F >= 5:
+				# Angular spectrum (p2p) propagation
+				_wf, x, y = self.angular_spectrum_propagation()
+
+			if output == 'complete':
+				_wf_list.append(_wf)
+				_x_list.append(x)
+				_y_list.append(y)
+		
+		if output == 'Complete':
+			return np.array([_wf_list, _x_list, _y_list])
+
+		else:
+			return np.array([_wf, x, y])
+
+		
+
+
+
+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.spiderTopHat(250, 100, 2001, 2001, dpix=1, oversample=1)
+	A = prof.topHat(2.5, n, m, dpix=dpix, oversample=1)
+
+	# Alternatively, create a gaussian profile
+	# gauss = GaussianBeam(0.3, l*1000)
+	# A = gauss.fundamentalMode(0,n,m,np.max(x))
+
+	# 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))
+	
+	# Start an optical train object. Add lenses and propagate. Currently not in use.
+	# LDE = OpticalTrain()
+
+	# 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.phaseDelayFunction(x,y,1.35e-6, 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.
+	oC = [-5, 5, -5, 5]
+	savename = 'simple_lens.pdf'
+	Exyz,_,_ = fou.FresnelPropagationFFTW(U_in, x, y, dz=L1.f, oversample=None, show=True, save=savename, outputCoordinates=oC)
+	
+	# 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 True:
+		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.FresnelPropagationFFTW(U_in, xin, yin, dz=i, oversample=None, 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(211)
+		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(212)
+		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)
+
+		# 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)))
+	
+
+

test cases (standalone)/glassCatalog.py

@@ -0,0 +1,12 @@
+#!/usr/bin/env Python3
+
+# Pyton dictionary with some glass types
+
+catalog = {
+	
+	'N-BK7' 	: [1.039612120E+00, 6.000698670E-03, 2.317923440E-01, 2.001791440E-02, 1.010469450E+00, 1.035606530E+02	],
+	'SiO2'		: [0.6961663, 		0.0684043, 		 0.4079426, 	  0.1162414, 	   0.8974794, 		9.896161		],
+	'F2'		: [1.345333590E+00, 9.977438710E-03, 2.090731760E-01, 4.704507670E-02, 9.373571620E-01, 1.118867640E+02	],
+	'S-FPL51'	: [1.170105050E+00, 6.162039240E-03, 4.757107830E-02, 2.633728760E-02, 7.638324450E-01, 1.418826420E+02	],
+	'S-LAH71'	: [1.982800310E+00, 1.189874560E-02, 3.167584500E-01, 5.271560010E-02, 2.444726460E+00, 2.132206970E+02	],
+}