From 61ca707689f66427ca0207960b6721bd4ade7634 Mon Sep 17 00:00:00 2001 From: yosata Date: Tue, 23 Jun 2020 16:05:59 +0200 Subject: [PATCH] Added features, fixed bugs and other small things --- .../generate_PSF_with_WFE.py | 20 + .../generate_bahtinov_dispersion.py | 69 + demonstration scripts/generate_field.py | 80 + .../generate_propagation_figures.py | 16 + .../thesis_propagation_figures.py | 51 + fourierOptics/__init__.py | 1 + fourierOptics/fourier.py | 36 +- fourierOptics/fourierOptics.py | 2095 ----------------- fourierOptics/micado.py | 438 ++++ fourierOptics/misc.py | 4 +- fourierOptics/optics.py | 43 +- fourierOptics/profiles.py | 49 +- 12 files changed, 784 insertions(+), 2118 deletions(-) create mode 100644 demonstration scripts/generate_PSF_with_WFE.py create mode 100644 demonstration scripts/generate_bahtinov_dispersion.py create mode 100644 demonstration scripts/generate_field.py create mode 100644 demonstration scripts/generate_propagation_figures.py create mode 100644 demonstration scripts/thesis_propagation_figures.py delete mode 100644 fourierOptics/fourierOptics.py create mode 100644 fourierOptics/micado.py diff --git a/demonstration scripts/generate_PSF_with_WFE.py b/demonstration scripts/generate_PSF_with_WFE.py new file mode 100644 index 0000000..745a763 --- /dev/null +++ b/demonstration scripts/generate_PSF_with_WFE.py @@ -0,0 +1,20 @@ +import numpy as np +import sys +sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics') +from fourierOptics import * + +if __name__ == '__main__': + # Initiate MICADO paraxial model + mcd = ParaxialMICADO() + + # Initiate WFE object + wfe = WavefrontError(mcd.n, mcd.m) + + # Generate PSD surfaces + psd = wfe.generate_height_map(200e-6, 2, 1) + + # Add it to ADC-P1 + mcd.ELT_ADC.p1.add_wavefront_error(psd) + + # Generate the PSF and plot it + E, x, y = mcd.propagate(l=1.35e-3,surface=-1, zoom=True, show=True) \ No newline at end of file diff --git a/demonstration scripts/generate_bahtinov_dispersion.py b/demonstration scripts/generate_bahtinov_dispersion.py new file mode 100644 index 0000000..dbae72b --- /dev/null +++ b/demonstration scripts/generate_bahtinov_dispersion.py @@ -0,0 +1,69 @@ +import numpy as np +import sys +sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics') +from fourierOptics import * + + +if __name__ == '__main__': + # Initiate MICADO paraxial model + mcd = ParaxialMICADO() + wfe = WavefrontError(mcd.n, mcd.m) + + # Create the Bahtinov mask and set it on the entrance pupil (ELT-M1) + bahtinov = Profiles().bahtinov_mask(40, 10, mcd.n, mcd.m, dpix=0.33, N=30, oversample=1) + mcd.ELT_CAM.add_mask(bahtinov) + + # Set up + Elist = [] + lList = np.linspace(1.49e-3, 1.78e-3, 15).tolist() + + # Dispersion of 40 mas over the band + dispMax = 40e-3 / 3600 # maximum dispersion in degrees + dispersion = np.linspace(0, dispMax, len(lList)).tolist() + + # Iterate over the wavelengths + for i in range(len(lList)): + l = lList[i] + + # Generating the tilt to the input field. + d = dispersion[i] + field_tilt = wfe.generate_tilt(mcd.x,mcd.y,tx=0, ty=d) + + # Adding the tilt to the input field + Ein = Profiles().ELT_pupil(mcd.n,mcd.m,mcd.dpix).astype('complex128')\ + * np.exp(-1j * 2*np.pi/l * field_tilt) + mcd.apply_mask(Ein, replace=True) + + # Propagate the input field throug the paraxial model + E, x, y = mcd.propagate(l=l,surface=-1, zoom=False, show=False, + image_output=[-4,4,-4,4], print_progress=True) + Elist.append(E) + + # Average over all the wavelengths to get the final field. + Eca = np.asarray([np.abs(x) for x in Elist]).mean(axis=0) # Amplitude + Eci = Eca**2 # Intensity + Ecp = np.asarray([np.angle(x) for x in Elist]).mean(axis=0) # Phase + + + # Convert the scale to sky-angles using the plate scale + x *= mcd.platescale*1000 # x in milli arcseconds + y *= mcd.platescale*1000 # y in milli arcseconds + + # Plotting the results + fig = plt.figure(figsize=(8,3)) + ax = fig.add_subplot(121) + im = ax.imshow(Eca, cmap='hot', origin='lower', extent=[x.min(),x.max(),y.min(),y.max()]) + ax.set_xlim([x.min(),x.max()]) + ax.set_ylim([y.min(),y.max()]) + ax.set_xlabel('x [mas]') + ax.set_ylabel('y [mas]') + fig.colorbar(im, ax=ax, label='Amplitude') + ax2 = fig.add_subplot(122) + im2 = ax2.imshow(Ecp, cmap='hot', origin='lower', extent=[x.min(),x.max(),y.min(),y.max()]) + ax2.set_xlim([x.min(),x.max()]) + ax2.set_ylim([y.min(),y.max()]) + ax2.set_xlabel('x [mas]') + ax2.set_ylabel('y [mas]') + fig.colorbar(im2, ax=ax2, label='Phase') + plt.tight_layout() + plt.show() \ No newline at end of file diff --git a/demonstration scripts/generate_field.py b/demonstration scripts/generate_field.py new file mode 100644 index 0000000..5eb02c4 --- /dev/null +++ b/demonstration scripts/generate_field.py @@ -0,0 +1,80 @@ +import numpy as np +import sys +sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics') +from fourierOptics import * +from pdb import set_trace + +def field_angle(hx, hy): + fx = hx * 0.0075 + fy = hy * 0.0075 + return fx, fy + + +if __name__ == '__main__': + # Initializing MICADO paraxial model + mcd = ParaxialMICADO() + wfe = WavefrontError(mcd.n, mcd.m) + + # Uncomment to load in WFE (coma) + # Problem here is that the field of view is too small to see + # a radial effect... + # xin, yin = load_variable('../adc_data_coordinates.pkl') + # adcObj = ADC_surfaces(xin,yin) + # zern = wfe.zernike_cartesian([0,0,0,0,0,0,0,100e-5], adcObj.xz, adcObj.yz, rescale=False) + # mcd.ELT_ADC.p1.add_wavefront_error(zern) + + # Setup + l = 1.49e-3 + Elist = [] + + # Set field points + hxList, hyList = np.mgrid[-0.03:0.03:3j,-0.03:0.03:3j] + hxList = hxList.flatten() + hyList = hyList.flatten() + + # Iterating over the field positions. + for i in range(len(hxList)): + hx = hxList[i] + hy = hyList[i] + + # Calculate the relative wavefront tilt for the field position + fieldx, fieldy = field_angle(hx,hy) + field_tilt = wfe.generate_tilt(mcd.x,mcd.y,tx=fieldx, ty=fieldy) + + # Apply the field tilt + Ein = Profiles().ELT_pupil(mcd.n,mcd.m,mcd.dpix).astype('complex128')\ + * np.exp(-1j * 2*np.pi/l * field_tilt) + mcd.apply_mask(Ein, replace=True) + + # Propagate through the MICADO system + E, x, y = mcd.propagate(l=l,surface=-1, zoom=False, show=False, + print_progress=False) + Elist.append(E) + + # Average over all the wavelengths to get the final field. + Eca = np.asarray([np.abs(x) for x in Elist]).mean(axis=0) # Amplitude + Eci = Eca**2 # Intensity + Ecp = np.asarray([np.angle(x) for x in Elist]).mean(axis=0) # Phase + + # Convert the scale to sky-angles using the plate scale + x *= mcd.platescale*1000 # x in milli arcseconds + y *= mcd.platescale*1000 # y in milli arcseconds + + # Plotting the results + fig = plt.figure(figsize=(8,3)) + ax = fig.add_subplot(121) + im = ax.imshow(Eca, cmap='hot', origin='lower', extent=[x.min(),x.max(),y.min(),y.max()]) + ax.set_xlim([x.min(),x.max()]) + ax.set_ylim([y.min(),y.max()]) + ax.set_xlabel('x [mas]') + ax.set_ylabel('y [mas]') + fig.colorbar(im, ax=ax, label='Amplitude') + ax2 = fig.add_subplot(122) + im2 = ax2.imshow(Ecp, cmap='hot', origin='lower', extent=[x.min(),x.max(),y.min(),y.max()]) + ax2.set_xlim([x.min(),x.max()]) + ax2.set_ylim([y.min(),y.max()]) + ax2.set_xlabel('x [mas]') + ax2.set_ylabel('y [mas]') + fig.colorbar(im2, ax=ax2, label='Phase') + plt.tight_layout() + plt.show() \ No newline at end of file diff --git a/demonstration scripts/generate_propagation_figures.py b/demonstration scripts/generate_propagation_figures.py new file mode 100644 index 0000000..922a09e --- /dev/null +++ b/demonstration scripts/generate_propagation_figures.py @@ -0,0 +1,16 @@ +import numpy as np +import sys +sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics') +from fourierOptics import * +import matplotlib.pyplot as plt +from pdb import set_trace + +if __name__ == '__main__': + plt.rc('font', family='serif') + plt.rc('axes', labelsize='large') + + # Initiate MICADO paraxial model + mcd = ParaxialMICADO() + + # Generate the PSF and plot it + E, x, y = mcd.propagate(l=1.35e-3,surface=-1, zoom=True, show=True, save='all') \ No newline at end of file diff --git a/demonstration scripts/thesis_propagation_figures.py b/demonstration scripts/thesis_propagation_figures.py new file mode 100644 index 0000000..0c88bb7 --- /dev/null +++ b/demonstration scripts/thesis_propagation_figures.py @@ -0,0 +1,51 @@ +import numpy as np +import sys +sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics') +from fourierOptics import * +import matplotlib.pyplot as plt +from pdb import set_trace + +if __name__ == '__main__': + plt.rc('font', family='serif') + plt.rc('axes', labelsize='large') + + # Initiate MICADO paraxial model + mcd = ParaxialMICADO(n=3001, m=3001, dpix=50) + + # Generate the PSF and plot it + # E, x, y = mcd.propagate(l=1.35e-3,surface=-1, zoom=True, show=True, save='all') + + confList = [ ( 1, [-0.4, 0.4, -0.4, 0.4], '02_fourier_ELT-IFP.pdf'), + ( 2, [-50, 50, -50, 50], '03_fourier_ELT-COL.pdf'), + ( 4, [-50, 50, -50, 50], '04_fourier_ELT-MCD-ADC-P2.pdf'), + (-1, [-1.2, 1.2, -1.2, 1.2], '05_fourier_ELT-MCD-DET.pdf')] + + + for i in range(len(confList)): + s = confList[i][0] + coordinates = confList[i][1] + savename = confList[i][2] + E, x, y = mcd.propagate(l=1.35e-3, surface=s, image_output=coordinates, + zoom=False, show=False, save=savename) + + Eout, x, y = mcd.propagate(l=1.35e-3, surface=0, zoom=False, show=False) + x /= 1e3 # convert to m + y /= 1e3 # convert to m + _fig = plt.figure(figsize=(8,3)) + _ax = _fig.add_subplot(121) + _im = _ax.imshow(np.abs(Eout), cmap='hot', extent=[x.min(), x.max(), y.min(), y.max()]) + _ax.set_xlim([-30, 30]) + _ax.set_ylim([-30, 30]) + _ax.set_xlabel('x [m]') + _ax.set_ylabel('y [m]') + _fig.colorbar(_im, ax=_ax, label='Amplitude') + _ax2 = _fig.add_subplot(122) + _im2 = _ax2.imshow(np.angle(Eout), cmap='hot', extent=[x.min(), x.max(), y.min(), y.max()]) + _ax2.set_xlim([-30, 30]) + _ax2.set_ylim([-30, 30]) + _ax2.set_xlabel('x [m]') + _ax2.set_ylabel('y [m]') + _fig.colorbar(_im2, ax=_ax2, label='Phase') + plt.tight_layout() + plt.savefig('01_fourier_ELT-M1.pdf', bbox_inches='tight', transparent=True) + plt.close()#show() \ No newline at end of file diff --git a/fourierOptics/__init__.py b/fourierOptics/__init__.py index 748f886..1a9800b 100644 --- a/fourierOptics/__init__.py +++ b/fourierOptics/__init__.py @@ -7,3 +7,4 @@ from .profiles import * from .psfFitter import * from .wavefrontErrors import * from .misc import * +from .micado import * \ No newline at end of file diff --git a/fourierOptics/fourier.py b/fourierOptics/fourier.py index 40a1af6..1f41227 100644 --- a/fourierOptics/fourier.py +++ b/fourierOptics/fourier.py @@ -1,5 +1,4 @@ import numpy as np -from pdb import set_trace import matplotlib.pyplot as plt from cmath import sqrt as csqrt from scipy import interpolate @@ -14,13 +13,14 @@ class FourierOptics: self.l = l # Wavelength - def plane_wave_propagation(self, A, xin, yin, dz=None, oversample=None, show=False, save=False, scaleInputArray=False, outputCoordinates=None): + def plane_wave_propagation(self, A, xin, yin, dz=None, oversample=None, show=False, + save=False, scaleInputArray=False, outputCoordinates=None): ''' Propagation through the Fresnel approximation, which is valid in many different situations. A separate code for the Fraunhofer propagation is not striclty necessary, as the exponent converges to 1 for large z values in any case. - Input parameters: + Parameters: A - [array] (Complex) Input field dpix - [float] pixel scale at which the calculations are done dz - [float] propagation distance @@ -32,7 +32,7 @@ class FourierOptics: outputCoordinates - [list] list the corners of the desired output coordinate plane. Order: [xmin, xmax, ymin, ymax] - Output: + Returns: E_out - [array] (Complex) Output field distribution x, y - [array] Output coordinates ''' @@ -59,9 +59,9 @@ class FourierOptics: _n, _m = np.shape(A) # Calculate the pixel size from the xi, eta plane - _dpix = np.mean(np.diff(xin)) + _dpix = np.mean(np.diff(xin.real)) if _dpix < 1e-99: - _dpix = np.mean(np.diff(xin)) + _dpix = np.mean(np.diff(xin.real)) # A at z=0 is then _E_xy0 = pyfftw.empty_aligned((_n,_m), dtype='complex128') @@ -93,8 +93,8 @@ class FourierOptics: # Coordinates of E_xy0 # _y, _x = np.mgrid[-_n / 2 * _dpix : _n / 2 * _dpix : 1j*_n, -_m / 2 * _dpix : _m / 2 * _dpix : 1j * _m] - _y, _x = yin, xin - + _y, _x = yin.real, xin.real + if outputCoordinates is not None: # Decide only use for plotting or can be used for output as well? # Performing interpolation to get to the required final plane, done for real and imaginary part. @@ -108,7 +108,6 @@ class FourierOptics: _E_out_ip_imag = _f_ip_imag(_xi, _yi) _E_xyz = _E_out_ip_real + _E_out_ip_imag*1j _y, _x = np.mgrid[outputCoordinates[0]:outputCoordinates[1]:1j*_n, outputCoordinates[2]:outputCoordinates[3]:1j * _m] - _E_xyz_extent = [np.min(_x.real), np.max(_x.real), np.min(_y.real), np.max(_x.real)] @@ -125,13 +124,13 @@ class FourierOptics: im2 = ax2.imshow(np.angle(_E_xyz), cmap='hot', extent=_E_xyz_extent) ax2.set_xlabel('x [mm]') ax2.set_ylabel('y [mm]') - fig.colorbar(im2, ax=ax2) + fig.colorbar(im2, ax=ax2, label='Phase') plt.tight_layout() if save: - plt.savefig(save, bbox_inches='tight') + plt.savefig(save, bbox_inches='tight', transparent=True) if show: plt.show() - elif not show: + else: plt.close() return np.array([_E_xyz, _x.real, _y.real]) @@ -144,7 +143,7 @@ class FourierOptics: A separate code for the Fraunhofer propagation is not striclty necessary, as the exponent converges to 1 for large z values in any case. - Input parameters: + Parameters: A - [array] (Complex) Input field dpix - [float] pixel scale at which the calculations are done dz - [float] propagation distance @@ -158,7 +157,7 @@ class FourierOptics: If 'auto': the output coordinates will be determined automatically. - Output: + Returns: E_out - [array] (Complex) Output field distribution x, y - [array] Output coordinates ''' @@ -189,7 +188,7 @@ class FourierOptics: _k = 2 * np.pi / self.l # Define (xi, eta) plane at z=0 in meters, with nxm points. We must use the pixel scale. - _xi, _eta = xin, yin + _xi, _eta = xin.real, yin.real # Calculate the pixel size from the xi, eta plane @@ -261,7 +260,7 @@ class FourierOptics: _xip, _yip = _x, _y - if show: + if show or save: # Plotting the results _fig = plt.figure(figsize=(8,3)) _ax = _fig.add_subplot(121) @@ -281,7 +280,10 @@ class FourierOptics: plt.tight_layout() if save: plt.savefig(save, bbox_inches='tight', transparent=True) - plt.show() + if show: + plt.show() + else: + plt.close() _x, _y = np.meshgrid(_xip, _yip) return np.array([_E_out_ip, _x.real, _y.real]) diff --git a/fourierOptics/fourierOptics.py b/fourierOptics/fourierOptics.py deleted file mode 100644 index 0071468..0000000 --- a/fourierOptics/fourierOptics.py +++ /dev/null @@ -1,2095 +0,0 @@ -#!/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/micado.py b/fourierOptics/micado.py new file mode 100644 index 0000000..4c69960 --- /dev/null +++ b/fourierOptics/micado.py @@ -0,0 +1,438 @@ +import numpy as np +from .fourier import * +from .gaussian import * +from .opticalTrain import * +from .optics import * +from .profiles import * +from .psfFitter import * +from .wavefrontErrors import * +from .misc import * + + + + +class ParaxialMICADO: + def __init__(self, l=1.35e-3, n=1001, m=1001, dpix=160): + self.l = l + self.dpix = dpix + self.n, self.m = n, m + self.x, self.y = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m)) + self.ELT_MASK = Profiles().ELT_pupil(n,m,dpix).astype('complex128') + + # Built optical system elements + self.ELT_MIRROR = ParabolicMirror(dx=0, dy=0, name='ELT-paraxial', CA=38542, f=6.84022e5) + self.ELT_FP = OpticalSurface(name='MCD-IFP') + self.ELT_COL = ParaxialLens(CA=154.54, f=1418.94, name='MCD-COL') + self.ELT_ADC = ADCObject() + self.ELT_CAM = ParaxialLens(CA=89.2, f=4278.32701, name='MCD-CAM') + self.ELT_DET = OpticalSurface(name='MCD-DET') + + # Built optical train + self.MCD = OpticalTrain() + self.MCD.add_object(self.ELT_MIRROR, dist=0) + self.MCD.add_object(self.ELT_FP, dist=self.ELT_MIRROR.f) + self.MCD.add_object(self.ELT_COL, dist=self.ELT_COL.f) + self.ELT_ADC.add_to_optical_train(self.MCD, 1209.4718) + self.MCD.add_object(self.ELT_CAM, dist=50) + self.MCD.add_object(self.ELT_DET, dist=self.ELT_CAM.f) + + self.last_PSF = None + self.clean_PSF = None + self.platescale = self.set_plate_scale() + print ('Your resolution is {}x{} with a pixel size of {}'.format(self.n,self.m,self.dpix)) + print('This means your grid is {} mm by {} mm, with each pixel {} mm.'.format( + self.x.max()-self.x.min(), self.y.max()-self.y.min(), self.dpix)) + + + def apply_mask(self, mask, replace=True): + ''' + Set a mask to the ELT entrance pupil. + ''' + if replace: + self.ELT_MASK = mask.astype('complex128') + elif not replace: + self.ELT_MASK *= mask + + + def set_PSF(self, PSF, x, y, perfect=True): + ''' + If calculating another PSF takes to long, + one can set a custom PSF using this + function. + + Can be set as the 'clean' PSF or as another + 'last' PSF, with the perfect flag. + ''' + if perfect: + self.clean_PSF = (PSF, x, y) + elif not perfect: + self.last_PSF = (PSF, x, y) + + + def centroid_shift(self, method='Coupling'): + ''' + Do a centroid shift calculation for the given method. + This is a quick way of doing it. It only returns the + centroid shift (and uncertainties) and no other parameters. + ''' + assert (self.last_PSF is not None), 'PSF necessary for fitting.' + + PSF, x, y = self.last_PSF + fitter = PSFFitter(x.real,y.real,np.abs(PSF)**2) + + if method == 'Coupling': + assert (self.clean_PSF is not None), 'Perfect PSF necessary for comparison.' + Eclean, xclean, yclean = self.clean_PSF + fitted_PSF, popt, var, coupling = fitter.fit_custom_field(Eclean, PSF, return_coupling=True) + x0, y0 = np.asarray(popt).T + sx0, sy0 = np.asarray(np.sqrt(var)).T + return x0, y0, sx0, sy0, coupling + + elif method == 'Airy': + fitted_PSF, popt, pcov = fitter.fit_airy_PSF(initialGuess=(0.03**2, 0.0, 0.0, 0.0, 0.03)) + _, x0, y0, _, _ = popt_airy + pcov_list = [ np.sqrt(pcov[i][i]) for i in range(len(pcov)) ] + _, sx0, sy0, _, _ = np.asarray(pcov_list).T + + elif method == 'Gaussian': + fitted_PSF, popt, pcov = fitter.fit_gaussian_PSF(initialGuess=(0.04, 0, 0, 0.01, 0.01, 0, 0)) + _, x0, y0, _, _, _, _ = popt + pcov_list = [ np.sqrt(pcov[i][i]) for i in range(len(pcov)) ] + _, sx0, sy0, _, _, _, _ = np.asarray(pcov_list).T + + elif method == 'Moffat': + fitted_PSF, popt, pcov = fitter.fit_moffat_PSF(initialGuess=(0.04, 0, 0, 0.01, 0.01, 0, 0, 4)) + _, x0, y0, _, _, _, _, _ = popt + pcov_list = [ np.sqrt(pcov[i][i]) for i in range(len(pcov)) ] + _, sx0, sy0, _, _, _, _, _ = np.asarray(pcov_list).T + + return x0, y0, sx0, sy0 + + + def propagate(self, l=None, surface=-1, zoom=False, load_up_to_adc=None, coordinates=None, + show=False, image_output=None, save=None, apply_interpolation=True, print_progress=True): + ''' + Propagate the input wavefront through the MICADO paraxial + model. Optionally the propagation starts at the ADC, if + a data file with the wavefront at that point is defined. + + Parameters: + l - Desired wavelength of the propagation + surface - Surface to propagate to. Default: -1 = Detector plane + zoom - Zoom in on the PSF on the detector? + load_up_to_adc - Use a .pkl file to load in the wavefront just before the adc. + coordinates - User defined coordinates that define the propagation sampling. + show_final - Plots the final amplitude and phase if True. + + Returns: + Exyz - Complex field with shape n,m describing the amplitude and phase + of the wavefront at the output surface. + x, y - Corresponding x and y coordinates, in mm. + ''' + # Set up wavelength + l = self.l if not l else l + + # Setting up the arguments for each surface propgation + surfaces = self.get_surfaces() + surface = len(surfaces)-1 if surface == -1 else surface + coordinatesList = self.get_coordinates(coordinates) + show,show_final_zoom= self.get_show_image(show, zoom, len(surfaces), surface) + output_coordinates = self.get_image_coordinates(image_output, coordinatesList, len(surfaces), surface) + saveList = self.get_save_image(save, len(surfaces), surface) + apply_ip = self.get_image_interpolation(apply_interpolation, zoom, len(surfaces), surface) + + # Set up the FourierOptics object for propagation + fou = FourierOptics(l=l) + x, y = self.x, self.y + p = 1 # Used as counter for the ADC prisms. + + # Load the data up to the ADC (useful for many iterations on the same + # calculation, where differences only appear at or after the ADC) + if load_up_to_adc is not None: + Eout, x, y = load_variable(load_up_to_adc) + x = x.real + y = y.real + s = 4 + elif load_up_to_adc is None: + s = 0 + + # Start propgation + if print_progress: print('Starting propagation...') + while s <= surface: + if (surfaces[s] == self.ELT_MASK).all(): + Ein = surfaces[s] + Eout = surfaces[s] + if show[s] or saveList[s]: + _fig = plt.figure(figsize=(8,3)) + _ax = _fig.add_subplot(121) + _im = _ax.imshow(np.abs(Eout), cmap='hot', extent=[x.min(), x.max(), y.min(), y.max()]) + _ax.set_xlim([x.min(), x.max()]) + _ax.set_ylim([y.min(), y.max()]) + _ax.set_xlabel('x [mm]') + _ax.set_ylabel('y [mm]') + _fig.colorbar(_im, ax=_ax, label='Amplitude') + _ax2 = _fig.add_subplot(122) + _im2 = _ax2.imshow(np.angle(Eout), cmap='hot', extent=[x.min(), x.max(), y.min(), y.max()]) + _ax2.set_xlim([x.min(), x.max()]) + _ax2.set_ylim([y.min(), y.max()]) + _ax2.set_xlabel('x [mm]') + _ax2.set_ylabel('y [mm]') + _fig.colorbar(_im2, ax=_ax2, label='Phase') + plt.tight_layout() + if saveList[s]: + plt.savefig(saveList[s], bbox_inches='tight', transparent=True) + if show[s]: + plt.show() + else: + plt.close() + + + elif surfaces[s] == self.ELT_MIRROR: + if print_progress: print('Propagating from ELT-M1 to ELT-IFP') + Ein = Eout * self.ELT_MIRROR.phase_delay(x,y,l,paraxial=True) + if s == surface and zoom: + xzoom, yzoom = x, y + Einzoom = Ein + fzoom = self.ELT_MIRROR.f + Eout, x, y = fou.fresnel_propagation(Ein, x, y, dz=self.ELT_MIRROR.f, oversample=None, + show=show[s], save=saveList[s], outputCoordinates=output_coordinates[s], + apply_interpolation=apply_ip[s]) + + elif surfaces[s] == self.ELT_FP: + if print_progress: print('Propagating from ELT-IFP to MCD-COL') + Ein = Eout + Eout, x, y = fou.fresnel_propagation(Ein, x, y, dz=self.ELT_COL.f, oversample=None, + show=show[s], save=saveList[s], outputCoordinates=output_coordinates[s], + apply_interpolation=apply_ip[s]) + + elif surfaces[s] == self.ELT_COL: + if print_progress: print('Propagating from MCD-COL to MCD-ADC-P1') + Ein = Eout * self.ELT_COL.phase_delay(x,y,l, paraxial=True) + Eout, x, y = fou.plane_wave_propagation(Ein, x, y, dz=self.MCD.sequentialDistance[3], oversample=None, + show=show[s], save=saveList[s], outputCoordinates=output_coordinates[s]) + + elif surfaces[s] == self.ELT_ADC and p == 1: + if print_progress: print('Propagating from MCD-ADC-P1 to MCD-ADC-P2') + Ein = Eout * self.ELT_ADC.get_phase_delay(x,y,l, theta=None, paraxial=True, prism=1) + Eout, x, y = fou.plane_wave_propagation(Ein, x, y, dz=self.MCD.sequentialDistance[4], oversample=None, + show=show[s], save=saveList[s], outputCoordinates=output_coordinates[s]) + p += 1 + + elif surfaces[s] == self.ELT_ADC and p == 2: + if print_progress: print('Propagating from MCD-ADC-P2 to MCD-ADC-P3') + Ein = Eout * self.ELT_ADC.get_phase_delay(x,y,l, theta=None, paraxial=True, prism=2) + Eout, x, y = fou.plane_wave_propagation(Ein, x, y, dz=self.MCD.sequentialDistance[5], oversample=None, + show=show[s], save=saveList[s], outputCoordinates=output_coordinates[s]) + p += 1 + + elif surfaces[s] == self.ELT_ADC and p == 3: + if print_progress: print('Propagating from MCD-ADC-P3 to MCD-ADC-P4') + Ein = Eout * self.ELT_ADC.get_phase_delay(x,y,l, theta=None, paraxial=True, prism=3) + Eout, x, y = fou.plane_wave_propagation(Ein, x, y, dz=self.MCD.sequentialDistance[6], oversample=None, + show=show[s], save=saveList[s], outputCoordinates=output_coordinates[s]) + p += 1 + + elif surfaces[s] == self.ELT_ADC and p == 4: + if print_progress: print('Propagating from MCD-ADC-P4 to MCD-CAM') + Ein = Eout * self.ELT_ADC.get_phase_delay(x,y,l, theta=None, paraxial=True, prism=4) + Eout, x, y = fou.plane_wave_propagation(Ein, x, y, dz=self.MCD.sequentialDistance[7], oversample=None, + show=show[s], save=saveList[s], outputCoordinates=output_coordinates[s]) + p += 1 + + elif surfaces[s] == self.ELT_DET: + if print_progress: print('Propagating from MCD-CAM to MCD-DET') + Ein = Eout * self.ELT_CAM.phase_delay(x,y,l, paraxial=True) + if s == surface and zoom: + xzoom, yzoom = x.real, y.real + Einzoom = Ein + fzoom = self.ELT_CAM.f + Eout, x, y = fou.fresnel_propagation(Ein, x, y, dz=self.ELT_CAM.f, oversample=None, + show=show[s], save=saveList[s], outputCoordinates=output_coordinates[s], apply_interpolation=apply_ip[s]) + s += 1 + + if zoom: + if print_progress: print('Zooming in on the PSF location...') + imax_1, imax_2 = find_maximum(np.abs(Eout)) + xmax = x[imax_1, imax_2] + ymax = y[imax_1, imax_2] + ind = s-3 if (s-3)>0 else 0 + newCoordinates = np.asarray(coordinatesList[ind]) + np.array([xmax, xmax, ymax, ymax]) + newCoordinates = newCoordinates.real.tolist() + Eout, x, y = fou.fresnel_propagation(Einzoom, xzoom, yzoom, dz=fzoom, oversample=None, + show=show_final_zoom, save=saveList[s-1], outputCoordinates=newCoordinates, apply_interpolation=True) + + if print_progress: print('Propagation finished!') + self.last_PSF = (Eout, x, y) + return Eout, x, y + + + def get_surfaces(self): + _surfaces = [ + self.ELT_MASK, + self.ELT_MIRROR, + self.ELT_FP, + self.ELT_COL, + self.ELT_ADC, + self.ELT_ADC, + self.ELT_ADC, + self.ELT_ADC, + self.ELT_CAM, + self.ELT_DET + ] + return _surfaces + + + def get_coordinates(self, coordinates): + # Output coordinates: + if coordinates is None: + _coordinatesList = [ + [-1.0, 1.0, -1.0, 1.0], + [-200, 200, -200, 200], + [-120, 120, -120, 120], + [-120, 120, -120, 120], + [-120, 120, -120, 120], + [-120, 120, -120, 120], + [-120, 120, -120, 120], + [-2.5, 2.5, -2.5, 2.5] + ] + elif coordinates is not None: + _coordinatesList = coordinates + return _coordinatesList + + def get_show_image(self, show, zoom, N, surf): + # Set show values. Will only show the image plane. + _showlist = [False] * N + _showlist[surf] = show + _show_final_zoom = False + if zoom: + _show_final_zoom = show + _showlist[surf] = False + _show = _showlist + return _show, _show_final_zoom + + def get_image_coordinates(self, image_output, coordinates, N, surf): + # Set the image coordinates + _output_coordinates = [None] * N + _output_coordinates[2] = coordinates[1] + if image_output is not None: + _output_coordinates[surf] = image_output + return _output_coordinates + + def get_save_image(self, save, N, surf): + # Set the save values + _savelist = [False] * N + if save == 'all': + _savelist = [ + '01_ELT_MASK_output.png', + '02_ELT_FP_output.png', + '03_ELT_MCD_COL_output.png', + '04_ELT_MCD_ADC_P1_output.png', + '05_ELT_MCD_ADC_P2_output.png', + '06_ELT_MCD_ADC_P3_output.png', + '07_ELT_MCD_ADC_P4_output.png', + '08_ELT_MCD_CAM_output.png', + '09_ELT_MCD_DET_output.png', + '10_ELT_MCD_DET_output_zoomed.png' + ] + elif save is not None: + _savelist[surf] = save + return _savelist + + def get_image_interpolation(self, apply_interpolation, zoom, N, surf): + # Apply interpolation to the output? + _apply_ip = [False] * N + if not zoom and apply_interpolation and surf!=2: + _apply_ip[surf] = True + return _apply_ip + + def set_plate_scale(self): + _ps = (3600*180/np.pi) / (self.ELT_MIRROR.CA * 53.4466) + self.platescale = _ps + return _ps + + + +class ADC_surfaces: + ''' + Very specific class to generate and scale wavefront errors on + the MICADO ADC prisms. + ''' + def __init__(self, x,y, clearAperture=125): + self.n, self.m = x.shape + self.clearAperture = clearAperture + self.radius = self.clearAperture / 2 + self.yz,self.xz = np.mgrid[x.min()/self.radius:x.max()/self.radius:self.n*1j, + x.min()/self.radius:x.max()/self.radius:self.m*1j] + self.zern_coefficients = None + self.zern_surface = None + self.psd_surface = None + self.latest_surface = None + + + def generate_ADC_wfe(self, zernRMS=20e-6, psdRMS=15e-6, zernModes=19): + ''' + We apply the wavefront error an each of the ADC prisms + as a phase transformation, corresponding to a + transmitted wavefront error. + + This wfe is separated into a zernike and psd component + with a total rms of 25 nm + zernike component - RMS: 20 nm + psd component - RMS: 15 nm + Total - RMS: 25 nm + ''' + print('Generating new ADC WFE...') + wfe = WavefrontError(self.n,self.m) + zern_rms = zernRMS # RMS of the zernike component of the WFE + midspat_rms = psdRMS # RMS of the PSD component of the WFE + + self.psd_surface = wfe.generate_height_map(midspat_rms, 2, 1) + zern_surface = wfe.generate_height_map(zern_rms, 2, 1) + zern_coefficients = wfe.fit_zernike_coefficients(zern_surface, modes=zernModes) + zern_coefficients[:3] = 0 # No defocus, tip or tilt + self.zern_coefficients = wfe.scale_zernike_coefficients(zern_coefficients, zern_rms) + + # Scale the zernike modes to the clear aperture of the ADC + self.zern_surface = wfe.zernike_cartesian(self.zern_coefficients, self.xz, self.yz, rescale=False) + + zern_surface = self.mask_surface(self.zern_surface)#, self.xz, self.yz) + psd_surface = self.mask_surface(self.psd_surface)#, self.xz, self.yz) + + self.latest_surface = zern_surface + psd_surface + return self.latest_surface + + + def mask_surface(self, A): + ''' + Set the array values to 0 where + the radius is > 1, without changing + the original array. + ''' + # Rescale x and y + _mask = np.sqrt(self.xz**2 + self.yz**2) > 1 + B = A.copy() + B[_mask] = 0 + return B + + + def scale_wfe(self, zernRMS=20e-5, psdRMS=15e-5): + ''' + Scaling of the original wfe, given the new + zernike and psd RMS values that are desired. + ''' + print('Scaling ADC WFE...') + _old_psd_rms = np.sqrt(np.mean(self.psd_surface**2)) + _factor_psd = psdRMS / _old_psd_rms + wfe = WavefrontError(self.n,self.m) + new_psd_surface = wfe.scale_rms(_factor_psd, data=self.psd_surface) + + new_zern_coefficients = wfe.scale_zernike_coefficients(self.zern_coefficients, zernRMS) + + # Scale the zernike modes to the clear aperture of the ADC + new_zern_surface = wfe.zernike_cartesian(new_zern_coefficients, self.xz, self.yz, rescale=False) + + new_zern_surface = self.mask_surface(new_zern_surface)#, _xz, _yz) + new_psd_surface = self.mask_surface(new_psd_surface)#, _xz, _yz) + + self.latest_surface = new_zern_surface + new_psd_surface + return self.latest_surface \ No newline at end of file diff --git a/fourierOptics/misc.py b/fourierOptics/misc.py index 7c90be2..162391a 100644 --- a/fourierOptics/misc.py +++ b/fourierOptics/misc.py @@ -90,6 +90,4 @@ def load_variable(filename): ''' with open(filename, 'rb') as f: variable = pickle.load(f) - return variable - - + return variable \ No newline at end of file diff --git a/fourierOptics/optics.py b/fourierOptics/optics.py index e1c0f48..fca9642 100644 --- a/fourierOptics/optics.py +++ b/fourierOptics/optics.py @@ -14,12 +14,15 @@ class LensObject: self.f = None self.sag = None self.wavefrontError = False + self.mask = None def add_wavefront_error(self, z): assert len(z.shape)==2, 'Wavefront data is not 2D. Cannot continue' self.sag = z self.wavefrontError = True + def add_mask(self, mask): + self.mask = mask @@ -42,6 +45,8 @@ class ParaxialLens(LensObject): if self.wavefrontError: _t_l *= np.exp(-1j * _k * self.sag / l) + if self.mask is not None: + _t_l *= self.mask # Masking out the parts that fall outside the lens clear aperture # dist_from_center = np.sqrt(x**2 + y**2) # _lensMask = (dist_from_center <= (self.CA/2)).astype(int) @@ -152,6 +157,9 @@ class SphericalLens(LensObject): if self.wavefrontError: _t_l *= np.exp(-1j * _k * (_n - 1) * self.sag) + if self.mask is not None: + _t_l *= self.mask + # Masking out the parts that fall outside the lens clear aperture dist_from_center = np.sqrt(x**2 + y**2) _lensMask = (dist_from_center <= (self.CA/2)).astype(int) @@ -291,6 +299,9 @@ class EvenAsphereLens(LensObject): if self.wavefrontError: _t_l *= np.exp(-1j * _k * (_n - 1) * self.sag) + if self.mask is not None: + _t_l *= self.mask + # Masking out the parts that fall outside the lens clear aperture dist_from_center = np.sqrt(x**2 + y**2) _lensMask = (dist_from_center <= (self.CA/2)).astype(int) @@ -317,13 +328,17 @@ class MirrorObject: self.type = None self.sag = None self.wavefrontError = False - + self.mask = None def add_wavefront_error(self, z): assert len(z.shape)==2, 'Wavefront data does not have the proper shape. Cannot continue' self.sag = z self.wavefrontError = True + def add_mask(self, mask): + self.mask = mask + + @@ -346,6 +361,9 @@ class ParabolicMirror(MirrorObject): else: _t_m = np.exp(1j * _k * 2 * _dz) + if self.mask is not None: + _t_m *= self.mask + # Masking out the parts that fall outside the mirror clear aperture dist_from_center = np.sqrt(x**2 + y**2) _mirrorMask = (dist_from_center <= (self.CA/2)).astype(int) @@ -375,6 +393,9 @@ class SphericalMirror(MirrorObject): else: _t_m = np.exp(1j * _k * 2 * _dz) + if self.mask is not None: + _t_m *= self.mask + # Masking out the parts that fall outside the mirror clear aperture dist_from_center = np.sqrt(x**2 + y**2) _mirrorMask = (dist_from_center <= (self.CA/2)).astype(int) @@ -400,6 +421,10 @@ class FreeformMirror(MirrorObject): _t_m = np.exp(1j * _k * 2 * ( self.z + self.sag) ) else: _t_m = np.exp(1j * _k * 2 * self.z) + + if self.mask is not None: + _t_m *= self.mask + return _t_m @@ -415,9 +440,17 @@ class OpticalSurface: if name: assert type(name)==str, 'No string for a name.' self.name = name + self.mask = None + + def add_mask(self, mask): + self.mask = mask def phase_delay(self, x, y, l, **kwargs): - return np.ones(np.shape(x)) + _t_l = np.ones(np.shape(x)) + + if self.mask is not None: + _t_l *= self.mask + return _t_l @@ -639,6 +672,7 @@ class Prism: self.sag2 = None self.wavefrontError = False self.thetaPrism = 0 + self.mask = None # Remember original tilts for reset: self.tilt1x_OG = tilt1 @@ -709,6 +743,9 @@ class Prism: self.sag2 = z self.wavefrontError = True + def add_mask(self, mask): + self.mask = mask + def phase_delay(self, x, y, l, theta=None, paraxial=True): ''' @@ -748,6 +785,8 @@ class Prism: elif self.sag2 is not None: _t_p *= np.exp(-1j * _k * (_n - 1) * self.sag2) + if self.mask is not None: + _t_p *= self.mask # _t_p *= _prismMask return _t_p diff --git a/fourierOptics/profiles.py b/fourierOptics/profiles.py index b1579ff..0fae11d 100644 --- a/fourierOptics/profiles.py +++ b/fourierOptics/profiles.py @@ -70,4 +70,51 @@ class Profiles: _x = np.linspace(-n/2*dpix, n/2*dpix, n) _y = np.linspace(-m/2*dpix, m/2*dpix, m) _ELT_pupil = _f_pupil(_x, _y) - return _ELT_pupil \ No newline at end of file + return _ELT_pupil + + + def bahtinov_mask(self, r1, r2, n, m, dpix, N, oversample=1, half_width_spiders=None): + ''' + Creates a Bahtinov mask profile. + This profile is a popular mask to focus a telescope in the + amateur community. + ''' + assert r1 > r2, "r2 < r1, change input" + + if half_width_spiders is None: + half_width_spiders = 0.03 * r1 + + # Obscured top hat + j,i = np.mgrid[-n/2*dpix : n/2*dpix : 1j*n*oversample, -m/2*dpix : m/2*dpix : 1j*m*oversample] + dist_from_center = np.sqrt(i**2 + j**2) + _Anew = (dist_from_center <= (r1)) & (dist_from_center >= (r2)) + _Anew = _Anew.astype(int) + + # Three quadrants of Bahtinov mask + _M1 = (i < -half_width_spiders) + _M2 = (i > half_width_spiders) & (j > half_width_spiders) + _M3 = (i > half_width_spiders) & (j < -half_width_spiders) + _MC = _M1 | _M2 | _M3 + + # Generating the lines + _d = 2 * r1 / int(N) + _M4 = step(j, i, period=_d, rotation=0) + _M5 = step(j, i, period=_d, rotation=30) + _M6 = step(j, i, period=_d, rotation=-30) + _MC = (_M1 & _M4) | (_M2 & _M5) | (_M3 & _M6) + + # Combining with the obscured top hat + _Anew = _Anew * _MC + return _Anew.astype(int) + + +def step(x, y, period=1, rotation=0): + ''' + Function that returns a 2D square wave with period and rotation. + ''' + theta = np.deg2rad(rotation) + xn = x * np.cos(theta) - y * np.sin(theta) + yn = x * np.sin(theta) + y * np.cos(theta) + a = ((2/period * xn ).astype(int) % 2) & (xn>= 0) + b = ((2/period * xn - 1).astype(int) % 2) & (xn< 0) + return a + b \ No newline at end of file