diff --git a/.gitignore b/.gitignore index b6e4761..cefc713 100644 --- a/.gitignore +++ b/.gitignore @@ -1,30 +1,6 @@ -# Byte-compiled / optimized / DLL files -__pycache__/ -*.py[cod] -*$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 +.DS_Store +__pychache__/ +.ipynb_checkpoints/ MANIFEST # PyInstaller diff --git a/fourierOptics/__init__.py b/fourierOptics/__init__.py new file mode 100644 index 0000000..748f886 --- /dev/null +++ b/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 * diff --git a/fourierOptics/data/ELT_M1_pupil_profile.fits b/fourierOptics/data/ELT_M1_pupil_profile.fits new file mode 100755 index 0000000..153e55e Binary files /dev/null and b/fourierOptics/data/ELT_M1_pupil_profile.fits differ diff --git a/fourierOptics/data/__init__.py b/fourierOptics/data/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/fourierOptics/data/__pycache__/__init__.cpython-37.pyc b/fourierOptics/data/__pycache__/__init__.cpython-37.pyc new file mode 100644 index 0000000..7308dd8 Binary files /dev/null and b/fourierOptics/data/__pycache__/__init__.cpython-37.pyc differ diff --git a/fourierOptics/data/__pycache__/glassCatalog.cpython-37.pyc b/fourierOptics/data/__pycache__/glassCatalog.cpython-37.pyc new file mode 100644 index 0000000..fe63863 Binary files /dev/null and b/fourierOptics/data/__pycache__/glassCatalog.cpython-37.pyc differ diff --git a/fourierOptics/data/glassCatalog.py b/fourierOptics/data/glassCatalog.py new file mode 100644 index 0000000..f50d65b --- /dev/null +++ b/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 ], +} \ No newline at end of file diff --git a/fourierOptics/fourier.py b/fourierOptics/fourier.py new file mode 100644 index 0000000..40a1af6 --- /dev/null +++ b/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] diff --git a/fourierOptics/fourierOptics.py b/fourierOptics/fourierOptics.py new file mode 100644 index 0000000..0071468 --- /dev/null +++ b/fourierOptics/fourierOptics.py @@ -0,0 +1,2095 @@ +#!/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 types +import time +from math import ceil +from scipy.special import factorial, j1 +import scipy.optimize as opt +from astropy.io import fits +import os + + + +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), 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 fresnel_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] + 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: + _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) + else: + _E_out_ip = _E_out + + + 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 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): + ''' + 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/#). + ''' + # 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=1, 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 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 + + + + + +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): + file_dir_path = os.path.dirname(os.path.realpath(__file__)) + fits_image_filename = file_dir_path + '/data/ELT_M1_pupil_profile.fits' + _ELT_pupil = fits.open(fits_image_filename)[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 + + + + + +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, 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] + + + + + +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 = 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): + if theta is not None: + self.set_orientation(theta) + + _t_p1 = self.p1.phase_delay(x, y, l, theta=self.theta, paraxial=False) + _t_p2 = self.p2.phase_delay(x, y, l, theta=self.theta, paraxial=False) + _t_p3 = self.p3.phase_delay(x, y, l, theta=180-self.theta, paraxial=False) + _t_p4 = self.p4.phase_delay(x, y, l, theta=180-self.theta, paraxial=False) + return _t_p1, _t_p2, _t_p3, _t_p4 + + + 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) + + + + +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.theta = 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): + # l in mm needs to be converted to l in micron + l *= 1e3 + 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_dispersion(l, a, b) + + + def calculate_dispersion(self, l, a, b): + # l in mm needs to be converted to l in micron + l *= 1e3 + 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.theta = 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 = x * np.tan(np.deg2rad(self.tilt1x)) + y * np.tan(np.deg2rad(self.tilt1y)) + self.second_surface_sag = x * np.tan(np.deg2rad(self.tilt2x)) + y * 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, 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=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? 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.theta + _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 + + + + + +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]) + + + + + +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): + ''' + 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 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): + ''' + 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) + 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) + 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 + + + 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 + + + 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 + + + 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 + + + + +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 + + + + + + + + + + + + diff --git a/fourierOptics/gaussian.py b/fourierOptics/gaussian.py new file mode 100644 index 0000000..9e79cda --- /dev/null +++ b/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] \ No newline at end of file diff --git a/fourierOptics/misc.py b/fourierOptics/misc.py new file mode 100644 index 0000000..7c90be2 --- /dev/null +++ b/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 + + diff --git a/fourierOptics/opticalTrain.py b/fourierOptics/opticalTrain.py new file mode 100644 index 0000000..3f123df --- /dev/null +++ b/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]) diff --git a/fourierOptics/optics.py b/fourierOptics/optics.py new file mode 100644 index 0000000..e1c0f48 --- /dev/null +++ b/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 + + \ No newline at end of file diff --git a/fourierOptics/profiles.py b/fourierOptics/profiles.py new file mode 100644 index 0000000..b1579ff --- /dev/null +++ b/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 \ No newline at end of file diff --git a/fourierOptics/psfFitter.py b/fourierOptics/psfFitter.py new file mode 100644 index 0000000..3f9aca6 --- /dev/null +++ b/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 diff --git a/fourierOptics/wavefrontErrors.py b/fourierOptics/wavefrontErrors.py new file mode 100644 index 0000000..3f3e25f --- /dev/null +++ b/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 diff --git a/test cases (standalone)/fourierOptics_01_diffracting_aperture.py b/test cases (standalone)/fourierOptics_01_diffracting_aperture.py new file mode 100644 index 0000000..1e2389e --- /dev/null +++ b/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))) + + + diff --git a/test cases (standalone)/fourierOptics_025_single_lens_system.py b/test cases (standalone)/fourierOptics_025_single_lens_system.py new file mode 100644 index 0000000..7631059 --- /dev/null +++ b/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))) + + + diff --git a/test cases (standalone)/fourierOptics_02_telescope.py b/test cases (standalone)/fourierOptics_02_telescope.py new file mode 100644 index 0000000..427c9d7 --- /dev/null +++ b/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))) + + + diff --git a/test cases (standalone)/fourierOptics_03_verification.py b/test cases (standalone)/fourierOptics_03_verification.py new file mode 100644 index 0000000..7f2a618 --- /dev/null +++ b/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() + + diff --git a/test cases (standalone)/fourierOptics_04_simple_lens_system.py b/test cases (standalone)/fourierOptics_04_simple_lens_system.py new file mode 100644 index 0000000..0b99f2d --- /dev/null +++ b/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))) + + + diff --git a/test cases (standalone)/fourierOptics_25_single_lens_system.py b/test cases (standalone)/fourierOptics_25_single_lens_system.py new file mode 100644 index 0000000..7631059 --- /dev/null +++ b/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))) + + + diff --git a/test cases (standalone)/glassCatalog.py b/test cases (standalone)/glassCatalog.py new file mode 100644 index 0000000..f50d65b --- /dev/null +++ b/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 ], +} \ No newline at end of file