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