#!/usr/bin/env Python3 import numpy as np from pdb import set_trace import matplotlib.pyplot as plt from cmath import sqrt as csqrt from scipy import interpolate from scipy.optimize import leastsq from matplotlib.ticker import FormatStrFormatter from matplotlib.colors import LogNorm import pyfftw import glassCatalog from tabulate import tabulate import time class FourierOptics: def __init__(self, l): self.l = l # Wavelength def planeWavePropagationFFTW(self, A, dpix=None, dz=None, oversample=None, show=False, save=False, scaleInputArray=False, outputLimits=False): ''' Propagation through the Fresnel approximation, which is valid in many different situations. A separate code for the Fraunhofer propagation is not striclty necessary, as the exponent converges to 1 for large z values in any case. Input parameters: A - [array] (Complex) Input field dpix - [float] pixel scale at which the calculations are done dz - [float] propagation distance oversample - [bool] Apply padding on the outside of A, increasing the size of A by the given factor. show - [bool] Show the resulting amplitude and phase map after propagation save - [bool] Save the figure scaleInputArray - [tuple] Rescale the resolution of the input array using interpolation. outputCoordinates - [list] list the corners of the desired output coordinate plane. Order: [xmin, xmax, ymin, ymax] Output: E_out - [array] (Complex) Output field distribution x, y - [array] Output coordinates ''' if scaleInputArray: A = np.asarray(A) _n, _m = A.shape _nA, _mA = np.linspace(-1, 1, _n), np.linspace(-1, 1, _m) A_ip_real = interpolate.interp2d(_nA, _mA, A.real, kind='cubic') A_ip_imag = interpolate.interp2d(_nA, _mA, A.imag, kind='linear') _a, _b = scaleInputArray _nx, _my = np.linspace(-1, 1, _a), np.linspace(-1, 1, _b) A_out_ip_real = A_ip_real(_nx, _my) A_out_ip_imag = A_ip_imag(_nx, _my) A_out_ip = A_out_ip_real + A_out_ip_imag*1j A = np.rint(A_out_ip_real) if oversample: _Aos = np.zeros((np.shape(A)[0]*oversample, np.shape(A)[1]*oversample)) _Aos[int(len(_Aos)/2-len(A)/2): int(len(_Aos)/2+len(A)/2), int(len(_Aos)/2-len(A)/2): int(len(_Aos)/2+len(A)/2)] = A A = _Aos _n, _m = np.shape(A) # A at z=0 is then _E_xy0 = pyfftw.empty_aligned((_n,_m), dtype='complex128') _E_xy0[:] = np.fft.ifftshift(A) # use ifftshift to shift origin to start array # _IntensityNormFactor = (np.abs(_E_xy0)**2).sum() # Factor to scale the intensity after the FFTs _IntensityNormFactor = np.sum(np.abs(_E_xy0)**2) _Az0 = pyfftw.empty_aligned((_n,_m), dtype='complex128') _FFTW_Obj = pyfftw.FFTW( _E_xy0, _Az0, axes=(0,1), direction='FFTW_BACKWARD', threads=6, flags=('FFTW_ESTIMATE',)) _FFTW_Obj() # E-field at z=dz is the fourier transform of A(z=0)*e^(-ikz) _fx_l = np.fft.fftfreq(_n, d=dpix) _fy_l = np.fft.fftfreq(_m, d=dpix) _fx, _fy = np.meshgrid(_fx_l, _fy_l) _kz = self.phaseFactorExponent(_fx, _fy, self.l) # origin at start of array _phaseFactor = np.exp(-1j * _kz * dz) # origin at start of array # Calculating the spatial frequencies in k-space _kx = np.fft.fftshift( _fx * 2 * np.pi ) _ky = np.fft.fftshift( _fy * 2 * np.pi ) _E_xyz = pyfftw.empty_aligned((_n,_m), dtype='complex128') _A_in = pyfftw.empty_aligned((_n,_m), dtype='complex128') _A_in[:] = _Az0 * _phaseFactor _FFTW_Obj2 = pyfftw.FFTW( _A_in, _E_xyz, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',)) _FFTW_Obj2() _E_xyz = np.fft.fftshift(_E_xyz) # use fftshift to move the origin from edge to center _E_xyz /= np.sqrt(_IntensityNormFactor) # Normalization back to initial energy. # Coordinates of E_xy0 _y, _x = np.mgrid[-_n / 2 * dpix : _n / 2 * dpix : 1j*_n, -_m / 2 * dpix : _m / 2 * dpix : 1j * _m] if outputLimits: # Decide only use for plotting or can be used for output as well? # Performing interpolation to get to the required final plane, done for real and imaginary part. _x = np.linspace(-_n/2*dpix, _n/2*dpix, _n) _y = np.linspace(-_m/2*dpix, _m/2*dpix, _m) _f_ip_real = interpolate.interp2d(_x, _y, _E_xyz.real, kind='linear') _f_ip_imag = interpolate.interp2d(_x, _y, _E_xyz.imag, kind='linear') _xi = np.linspace(outputLimits[0],outputLimits[1],_n) _yi = np.linspace(outputLimits[0],outputLimits[1],_m) _E_out_ip_real = _f_ip_real(_xi, _yi) _E_out_ip_imag = _f_ip_imag(_xi, _yi) _E_xyz = _E_out_ip_real + _E_out_ip_imag*1j _y, _x = np.mgrid[outputLimits[0]:outputLimits[1]:1j*_n, outputLimits[0]:outputLimits[1]:1j * _m] _E_xyz_extent = [np.min(_x), np.max(_x), np.min(_y), np.max(_x)] if show or save: fig = plt.figure(figsize=(8,3)) ax = fig.add_subplot(121) im = ax.imshow(np.abs(_E_xyz), cmap='hot', extent=_E_xyz_extent)#, vmax=1, vmin=0) ax.set_xlabel('x [mm]') ax.set_ylabel('y [mm]') fig.colorbar(im, ax=ax, label='Normalized intensity') ax2 = fig.add_subplot(122) im2 = ax2.imshow(np.angle(_E_xyz), cmap='hot', extent=_E_xyz_extent) ax2.set_xlabel('x [mm]') ax2.set_ylabel('y [mm]') fig.colorbar(im2, ax=ax2) plt.tight_layout() if save: plt.savefig(save, bbox_inches='tight') if show: plt.show() elif not show: plt.close() return np.array([_E_xyz, _x, _y]) def FresnelPropagationFFTW(self, A, xin, yin, dz=None, oversample=None, show=False, save=False, scaleInputArray=False, outputCoordinates='auto'): ''' Propagation through the Fresnel approximation, which is valid in many different situations. A separate code for the Fraunhofer propagation is not striclty necessary, as the exponent converges to 1 for large z values in any case. Input parameters: A - [array] (Complex) Input field dpix - [float] pixel scale at which the calculations are done dz - [float] propagation distance oversample - [bool] Apply padding on the outside of A, increasing the size of A by the given factor. show - [bool] Show the resulting amplitude and phase map after propagation save - [bool] Save the figure scaleInputArray - [tuple] Rescale the resolution of the input array using interpolation. outputCoordinates - [list] list the corners of the desired output coordinate plane. Order: [xmin, xmax, ymin, ymax] If 'auto': the output coordinates will be determined automatically. Output: E_out - [array] (Complex) Output field distribution x, y - [array] Output coordinates ''' assert np.shape(xin) == np.shape(A), 'Input coordinates (xin) do not match shape of input field.' assert np.shape(yin) == np.shape(A), 'Input coordinates (yin) do not match shape of input field.' A = np.conj(A) if scaleInputArray: A = np.asarray(A) _n, _m = A.shape _nA, _mA = np.linspace(-1, 1, _n), np.linspace(-1, 1, _m) A_ip_real = interpolate.interp2d(_nA, _mA, A.real, kind='linear') A_ip_imag = interpolate.interp2d(_nA, _mA, A.imag, kind='linear') _a, _b = scaleInputArray _nx, _my = np.linspace(-1, 1, _a), np.linspace(-1, 1, _b) A_out_ip_real = A_ip_real(_nx, _my) A_out_ip_imag = A_ip_imag(_nx, _my) A_out_ip = A_out_ip_real + A_out_ip_imag*1j A = np.rint(A_out_ip_real) # if oversample: # _Aos = np.zeros((np.shape(A)[0]*oversample, np.shape(A)[1]*oversample)) # _Aos[int(len(_Aos)/2-len(A)/2): int(len(_Aos)/2+len(A)/2), int(len(_Aos)/2-len(A)/2): int(len(_Aos)/2+len(A)/2)] = A # A = _Aos _n, _m = np.shape(A) _k = 2 * np.pi / self.l # Define (xi, eta) plane at z=0 in meters, with nxm points. We must use the pixel scale. _xi, _eta = xin, yin # Calculate the pixel size from the xi, eta plane _dpix = np.mean(np.diff(_xi)) if _dpix < 1e-99: _dpix = np.mean(np.diff(_xi.T)) # Gaussian test: use 8*w0 radius for FFT: if outputCoordinates == 'auto': _w0_fit, _ = leastsq(self.minimize, self.l, args=(0, _xi, _eta, A)) _gauss = GaussianBeam(_w0_fit, self.l) _halfwidth = 8 * _gauss.beamRadius(dz)[0] _xip = np.linspace(-_halfwidth, _halfwidth, _n) # Desired x,y - plane _yip = np.linspace(-_halfwidth, _halfwidth, _m) # Desired x,y - plane else: _xip = np.linspace(outputCoordinates[0], outputCoordinates[1], _n) # Desired x,y - plane _yip = np.linspace(outputCoordinates[2], outputCoordinates[3], _m) # Desired x,y - plane # The quadratic phase exponential of the input wavefront _phaseComponent = np.exp(1j * _k / (2 * dz) * (_xi**2 + _eta**2)) # Complete input _U_in = pyfftw.empty_aligned((_n,_m), dtype='complex128') _U_in[:] = A * _phaseComponent # Output of the FFT evaluated at frequency fx = x / (lambda * z) _U_out = pyfftw.empty_aligned((_n,_m), dtype='complex128') _FFTW_Obj = pyfftw.FFTW( np.fft.ifftshift(_U_in), _U_out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',)) _FFTW_Obj() _U_out = np.fft.fftshift(_U_out) # Evaluate frequencies: _fx = np.fft.fftshift(np.fft.fftfreq(np.shape(_xi)[0], d=_dpix)) _fy = np.fft.fftshift(np.fft.fftfreq(np.shape(_xi)[1], d=_dpix)) _x = _fx * self.l * dz _y = _fy * self.l * dz _x2, _y2 = np.meshgrid(_x, _y) # Final output and frequencies _pre_amble = np.exp(1j * _k * dz) / (1j * self.l * dz) * np.exp(1j * _k / (2 * dz) * (_x2**2 + _y2**2)) _E_out = _pre_amble * _U_out # Normalization of the output values _IntensityNormFactor = np.sum(np.abs(A)**2) _E_out /= (np.abs(np.exp(1j*_k*dz)/(1j*self.l*dz)) * np.sqrt(_IntensityNormFactor) * np.sqrt(_n*_m)) # Performing interpolation to get to the required final plane, done for real and imaginary part. _f_ip_real = interpolate.interp2d(_x, _y, _E_out.real, kind='linear') _f_ip_imag = interpolate.interp2d(_x, _y, _E_out.imag, kind='linear') _E_out_ip_real = _f_ip_real(_xip, _yip) _E_out_ip_imag = _f_ip_imag(_xip, _yip) _E_out_ip = _E_out_ip_real + _E_out_ip_imag*1j _E_out_ip /= np.sqrt(((_x2.max()/_xip.max())*(_y2.max()/_yip.max()))) _E_out_ip = np.conj(_E_out_ip) if show: # Plotting the results _fig = plt.figure(figsize=(8,3)) _ax = _fig.add_subplot(121) # _im = _ax.imshow(np.sqrt(np.abs(_E_out_ip)), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()]) _im = _ax.imshow(np.abs(_E_out_ip), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()]) _ax.set_xlim([_xip.min(), _xip.max()]) _ax.set_ylim([_yip.min(), _yip.max()]) _ax.set_xlabel('x [mm]') _ax.set_ylabel('y [mm]') _fig.colorbar(_im, ax=_ax, label='Amplitude') # _fig.colorbar(_im, ax=_ax, label='$\\sqrt{Amplitude}$') _ax2 = _fig.add_subplot(122) _im2 = _ax2.imshow(np.angle(_E_out_ip), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()]) _ax2.set_xlim([_xip.min(), _xip.max()]) _ax2.set_ylim([_yip.min(), _yip.max()]) _ax2.set_xlabel('x [mm]') _ax2.set_ylabel('y [mm]') _fig.colorbar(_im2, ax=_ax2, label='Phase') plt.tight_layout() if save: plt.savefig(save, bbox_inches='tight', transparent=True) plt.show() _x, _y = np.meshgrid(_xip, _yip) return np.array([_E_out_ip, _x, _y]) def freqArray(self, A, dpix): n, m = np.shape(A) fy, fx = np.mgrid[0:m, 0:n] fx = fx.flatten() fy = fy.flatten() for i in range(len(fx)): fx[i] = (fx[i] - n*bool(fx[i]>(n-1)/2) ) / (n*dpix) fy[i] = (fy[i] - m*bool(fy[i]>(m-1)/2) ) / (m*dpix) fx = fx.reshape(n,m) fy = fy.reshape(n,m) return fx, fy def phaseFactorExponent(self, fx, fy, l): k = 2*np.pi / l kx = np.asarray(2*np.pi * fx) ky = np.asarray(2*np.pi * fy) K_squared = (kx**2 + ky**2) kA = np.ones(K_squared.shape) * k K_1 = np.sqrt((kA**2-K_squared).astype('complex128')) kz = K_1.real + -1j*K_1.imag return kz def minimize(self, a, dz, x, y, og_data): ''' Residual function used for fitting the gaussian beam to the input distribution. Input: a - [float] Guess for w0 dz - [float] Propagation distance, should be known x,y - [array] meshgrid of the x and y coordinates. og_data - [array] input distribution to be fitted. ''' return np.abs((self.gaussianBeamFit(a,dz,x,y) - og_data).flatten()).real def gaussianBeamFit(self, a, dz, x, y): ''' Function used to fit the Gaussian beam to an input plane. Very similar to the GaussianBeam class below, but modified to allow a best fit. Input: a - [float] Guess for w0 dz - [float] Propagation distance, should be known x,y - [array] meshgrid of the x and y coordinates. ''' _w0 = a _z = dz _r = np.sqrt(x**2 + y**2) if _z != 0: _z_c = np.pi * _w0**2 / self.l _w = _w0 * np.sqrt(1 + (_z/_z_c)**2) _R = _z + _z_c **2 / _z _phi0 = np.arctan(_z/_z_c) _k = 2 * np.pi / self.l _results = np.sqrt(2 / (np.pi*_w**2)) * np.exp(- _r**2 / _w**2 - 1j * _k * _z - 1j * np.pi * _r**2 / (self.l * _R) + 1j*_phi0) elif _z == 0: _results = np.sqrt(2 / (np.pi*_w0**2)) * np.exp( - _r**2 / _w0**2 ) return _results class profiles: def __init__(self): pass def topHat(self, r, n, m, dpix, oversample=1): j,i = np.mgrid[-n/2*dpix : n/2*dpix : 1j*n*oversample, -m/2*dpix : m/2*dpix : 1j*m*oversample] dist_from_center = np.sqrt(i**2 + j**2) _Anew = (dist_from_center <= (r)).astype(int) return _Anew def obscuredTopHat(self, r1, r2, n, m, dpix, oversample=1): ''' Creates a grid of n x m points, following a top hat function where all points inside a radius r is 1 and 0 otherwise. ''' assert r1 > r2, "r2 < r1, change input" j,i = np.mgrid[-n/2*dpix : n/2*dpix : 1j*n*oversample, -m/2*dpix : m/2*dpix : 1j*m*oversample] dist_from_center = np.sqrt(i**2 + j**2) _Anew = (dist_from_center <= (r1)) & (dist_from_center >= (r2)) return _Anew.astype(int) def spiderTopHat(self, r1, r2, n, m, dpix, oversample=1): ''' Creates a grid of n x m points, following a top hat function where all points inside a radius r is 1 and 0 otherwise. ''' assert r1 > r2, "r2 < r1, change input" j,i = np.mgrid[-n/2*dpix : n/2*dpix : 1j*n*oversample, -m/2*dpix : m/2*dpix : 1j*m*oversample] n *= oversample m *= oversample dist_from_center = np.sqrt(i**2 + j**2) _Anew = (dist_from_center <= (r1)) & (dist_from_center >= (r2)) _Aspider = (np.abs(i) < (0.05*r1)) | (np.abs(j) < (0.05*r1)) _Anew = _Anew & ~_Aspider return _Anew class GaussianBeam: def __init__(self, w0, l): # Confocal distance self.z_c = np.pi * w0**2 / l # Far-field divergence angle self.theta_0 = l / (np.pi * w0) self.l = l self.w0 = w0 self.fwhm = 1.18 * self.theta_0 def fundamentalMode(self, z, n, m, rmax): y, x = np.mgrid[-rmax:rmax:n*1j,-rmax:rmax:n*1j] r = np.sqrt(x**2 + y**2) if z != 0: _w = self.beamRadius(z) _R = self.curvature(z) _phi0 = self.phaseShift(z) _k = 2 * np.pi / self.l return np.sqrt(2 / (np.pi*_w**2)) * np.exp(- r**2 / _w**2 - 1j * _k * z - 1j * np.pi * r**2 / (self.l * _R) + 1j*_phi0) elif z == 0: return np.sqrt(2 / (np.pi*self.w0**2)) * np.exp( - r**2 / self.w0**2 ) def curvature(self, z): return z + self.z_c **2 / z def beamRadius(self, z): return self.w0 * np.sqrt(1 + (z/self.z_c)**2) def phaseShift(self, z): return np.arctan(z/self.z_c) def relativePower(self, r, z): return np.exp(-2 * (r / self.beamRadius(z))**2) class lensObject: ''' The lens object describes any spherical lens and find the phase delay for all light traveling through it. The refractive index, of course, depends on the wavelength. ''' def __init__(self, glass, R1, R2, thickness, CA, f=None, name=None): self.glass = glass assert glassCatalog.catalog.get(self.glass), 'Glass type not in catalog' self.R1 = R1 # mm self.R2 = R2 # mm self.thickness = thickness # mm self.CA = CA # mm self.dispersionCoefficients = glassCatalog.catalog.get(self.glass) self.referenceWavelength = 1.35 # micron self.type = 'Spherical lens' if not f: self.f = self.findFocalLength() else: self.f = f if name: assert type(name)==str, 'No string for a name.' self.name = name else: self.name = str(self) def findRefractiveIndex(self, _l): _a = self.dispersionCoefficients _n = np.sqrt( (_a[0]* _l**2)/(_l**2 - _a[1]) + (_a[2] * _l**2) / (_l**2 - _a[3]) + (_a[4] * _l**2) / (_l**2 - _a[5]) + 1) return _n def findFocalLength(self): _n = self.findRefractiveIndex(self.referenceWavelength) _f_recip = (_n - 1) * (1 / self.R1 - 1 / self.R2) return 1/_f_recip def phaseDelayFunction(self, x, y, l, paraxial=False): if not paraxial: _A1 = self.R1 * (1 - np.sqrt(1 - (x**2 + y**2) / self.R1**2 ) ) _A2 = self.R2 * (1 - np.sqrt(1 - (x**2 + y**2) / self.R2**2 ) ) _dxy = self.thickness - _A1 + _A2 if paraxial: _dxy = self.thickness - (x**2 + y**2)/2 * (1/self.R1 - 1/self.R2) _k = 2 * np.pi / 30e-3 _n = self.findRefractiveIndex(self.referenceWavelength) _t_l = np.exp(-1j * _k * self.thickness)* np.exp(-1j * _k * (_n - 1) * _dxy) # Masking out the parts that fall outside the lens clear aperture dist_from_center = np.sqrt(x**2 + y**2) _lensMask = (dist_from_center <= (self.CA)).astype(int) _t_l *= _lensMask return _t_l class mirrorObject: ''' Describes any mirror. No wavelength dependence exist. Coatings are not considered. ''' def __init__(self, dx, dy, name=None): self.dx = dx self.dy = dy if name: assert type(name)==str, 'No string for a name.' self.name = name else: self.name = str(self) class parabolicMirror(mirrorObject): def __init__(self, f, CA, dx=None, dy=None, name=None): mirrorObject.__init__(self, dx=dx, dy=dy, name=name) self.f = f self.CA = CA self.type = 'Parabolic mirror' def phaseDelayFunction(self, x, y, l): _k = 2 * np.pi / l _dz = (x**2 + y**2) / (4 * self.f) _t_m = np.exp(1j * _k * 2 * _dz) return _t_m class sphericalMirror(mirrorObject): def __init__(self, R, CA, dx=None, dy=None, name=None): mirrorObject.__init__(self, dx=dx, dy=dy, name=name) self.R = R self.CA = CA self.type = 'Spherical mirror' def phaseDelayFunction(self, x, y, l): # _k = 2 * np.pi / l # _dz = self.R * (1 - np.sqrt(1 - (x**2 + y**2)/self.R**2) ) # _t_m = np.exp(1j * _k * 2 * _dz) # return _t_m _k = 2 * np.pi / l _dz = (x**2 + y**2) / (4 * self.f) # Masking out the parts that fall outside the mirror clear aperture dist_from_center = np.sqrt(x**2 + y**2) _mirrorMask = (dist_from_center <= (self.CA/2)).astype(int) _t_m *= _mirrorMask return _t_m class freeformMirror(mirrorObject): def __init__(self, x, y, z, dx=None, dy=None, name=None): mirrorObject.__init__(self, dx=dx, dy=dy, name=name) self.x = x self.y = y self.z = z self.type = 'User def. mirror' def phaseDelayFunction(self, l): _k = 2 * np.pi / l _t_m = np.exp(1j * _k * 2 * self.z) return _t_m class OpticalTrain: ''' The main class where an optical system is build up sequentially. ''' def __init__(self): self.opticalTrain = [] self.sequentialDistance = [] def add_object(self, obj, dist): ''' Adding a lens or mirror object to the opticalTrain obj The lens or mirror object dist The distance from the the previous object. ''' assert type(dist) == float or type(dist) == int, 'Distance is not a float or integer.' self.opticalTrain.append(obj) self.sequentialDistance.append(float(dist)) def print_lens_Data(self): ''' Print the data of the current opticalTrain. ''' _tableData = [] for _lens in self.opticalTrain: _i = self.opticalTrain.index(_lens) _distances = self.sequentialDistance _distances.append('-') _tableData.append([_i, _lens.name, _lens.type, _distances[_i+1] ]) print(tabulate(_tableData, headers=['#', 'Name', 'Type', 'Distance to next surface'])) def propagate_system(self, U_in, x, y, l, output=None, paraxial=True): # We need to have an input distribution # We need to determine what propagation method is required # Call the objects in the self.opticalTrain sequentially # Call the subsequent phaseDelay if output == 'complete': _wf_list = [] _x_list = [] _y_list = [] _wf_list.append(_wf) _x_list.append(x) _y_list.append(y) _wf = U_in for _lens in self.opticalTrain: _i = self.opticalTrain.index(_lens) _dz = self.sequentialDistance[_i] _wf = _lens.phaseDelayFunction(x, y, l, paraxial=paraxial) # Perform Fresnel number test: _F = self.fresnel_number() if _F < 5: # Fresnel propagation _wf, x, y = self.fresnel_propagation() elif _F >= 5: # Angular spectrum (p2p) propagation _wf, x, y = self.angular_spectrum_propagation() if output == 'complete': _wf_list.append(_wf) _x_list.append(x) _y_list.append(y) if output == 'Complete': return np.array([_wf_list, _x_list, _y_list]) else: return np.array([_wf, x, y]) if __name__ == '__main__': plt.rc('font', family='serif') plt.rc('axes', labelsize='large') # Start time counter start = time.time() # Pick a wavelength l = 30e-3 # Initiate Fourier Optics fou = FourierOptics(l=l) # Pick sampling parameters dpix = 8 n, m = 4001, 4001 print ('Your resolution is {}x{} with a pixel size of {}'.format(n,m,dpix)) # Initiate the profiles class and pick a profile (tophat, obscured tophat or spider) prof = profiles() A = prof.topHat(4000, n, m, dpix=dpix, oversample=1) # A = prof.spiderTopHat(4000, 1000, n, m, dpix=dpix, oversample=1) # A = prof.obscuredTopHat(4000, 1000, n, m, dpix=dpix, oversample=1) # Alternatively, create a gaussian profile # gauss = GaussianBeam(0.3, l*1000) # A = gauss.fundamentalMode(0,n,m,np.max(x)) # Create an input grid x, y = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m)) # Start an optical train object. Add lenses and propagate. Currently not in use. # LDE = OpticalTrain() # Create a mirror with a focal length 64 m, with a clear aperture radius of 4 m. M1 = parabolicMirror(dx=0, dy=0, name='M1', CA=8000, f=64000) print('Focal length of mirror {} is {:.1f} mm.'.format(M1.name, M1.f)) # Calculate the phase delay introduced by the mirror t_l = M1.phaseDelayFunction(x,y,l) # The input distribution gets a curved wavefront after the mirror U_in = A * t_l plt.imshow(np.abs(U_in)) plt.colorbar() plt.savefig('telescope_mask.pdf', bbox_inches='tight', transparent=True) plt.show(block=False) # Performing the propagation using Fresnel approximation to the focal plane. # oC = [-0.001, 0.001, -0.001, 0.001] oC = [-2, 2, -2, 2] # savename = 'fourierOptics_02_telescope_fresnel.pdf' Exyz,xout,yout = fou.FresnelPropagationFFTW(U_in, x, y, dz=M1.f, oversample=None, show=True, save=False, outputCoordinates=oC) # Plot a cross section of the new field (useful to better see the diffraction rings) fig = plt.figure() ax = fig.add_subplot(111) ax.grid(True, alpha=0.3) ax.semilogy(np.linspace(np.min(oC), np.max(oC),n), np.abs(Exyz[int((n-1)/2),:]**2), 'r-') ax.set_xlabel('x (mm)') ax.set_ylabel('Intensity') ax.tick_params(axis ='both', direction='in', which='both') # plt.savefig('fourierOptics_02_telescope_slice.pdf', bbox_inches='tight', transparent=True) plt.show() # Combined plot # Plotting the results _fig = plt.figure(figsize=(12,3)) _ax = _fig.add_subplot(131) _im = _ax.imshow(np.abs(Exyz), cmap='hot', extent=oC) _ax.set_xlim([oC[0], oC[1]]) _ax.set_ylim([oC[2], oC[3]]) _ax.set_xlabel('x [mm]') _ax.set_ylabel('y [mm]') _fig.colorbar(_im, ax=_ax, label='Amplitude') _ax2 = _fig.add_subplot(132) _im2 = _ax2.imshow(np.angle(Exyz), cmap='hot', extent=oC) _ax2.set_xlim([oC[0], oC[1]]) _ax2.set_ylim([oC[2], oC[3]]) _ax2.set_xlabel('x [mm]') _ax2.set_ylabel('y [mm]') _fig.colorbar(_im2, ax=_ax2, label='Phase') _ax3 = _fig.add_subplot(133) _ax3.grid(True, alpha=0.3) _ax3.semilogy(np.linspace(np.min(oC), np.max(oC),n), np.abs(Exyz[int((n-1)/2),:]**2), 'r-') _ax3.set_xlabel('x (mm)') _ax3.set_ylabel('Intensity') _ax3.tick_params(axis ='both', direction='in', which='both') plt.tight_layout() plt.savefig('fourierOptics_02_telescope_complete.pdf', bbox_inches='tight', transparent=True) plt.show() # The following produces a cross-section of the propagation for various z, so that you can see how the beam evolves. if False: Ics = [] Ics_x = [] Isum = [] xmax = [] zList = np.linspace(20,220,20).tolist() xin, yin = x, y for i in zList: title = f"{zList.index(i)+1:03d}" print (' Calculating {0}/{1}'.format(title, f"{len(zList):03d}"), end="\r") # Propagate the wavefront # Exyz, x, y = fou.planeWavePropagationFFTW(A, dz=i*l, dpix=dpix, oversample=1, show=False, save=False)#, scaleInputArray=(n,m)) Exyz, x, y = fou.FresnelPropagationFFTW(U_in, xin, yin, dz=i, oversample=None, show=False, save=False, outputCoordinates=oC) # Process the results by taking a slice of the data Isum.append(np.sum(np.abs(Exyz)**2)) Ics.append((np.abs(Exyz[int(Exyz.shape[0]/2)])**2)) Ics_x.append(x.real[0]) y_ext = [y.real[0,0],y.real[-1,0]] # Reset the data x = None y = None # Plotting all slices fig = plt.figure(figsize=(8,4)) ax = fig.add_subplot(211) ax.tick_params(axis ='both', direction='in') # Plotting the square root of the amplitude for better contrast im = ax.imshow(np.sqrt(np.asarray(Ics)).T, cmap='hot', aspect='auto', interpolation='bicubic', extent=[zList[0],zList[-1],y_ext[0],y_ext[1]]) ax.yaxis.set_major_formatter(FormatStrFormatter('%.1E')) # Plotting the intensity relative to the input ax2 = fig.add_subplot(212) ax2.plot(zList, np.asarray(Ics)[:,int(np.asarray(Ics).shape[1]/2)], 'k-') ax2.yaxis.set_major_formatter(FormatStrFormatter('%.1E')) ax2.set_xlabel('Propagation distance (mm)') ax2.tick_params(axis ='both', direction='in') ax2.set_ylabel('Intensity') ax2.set_xlim([np.min(zList), np.max(zList)]) ax2.grid(True, alpha=0.3) # Saving the figure? Then show. plt.tight_layout() # plt.savefig('fresnel_cross.pdf', bbox_inches='tight') # plt.savefig('fresnel_cross.png', bbox_inches='tight') plt.show() # Closing down the script and stopping the timer. end = time.time() print('The program took {:.3f} seconds!'.format((end-start)))