#!/usr/bin/env Python3 import numpy as np from pdb import set_trace import matplotlib.pyplot as plt from matplotlib import cm, colors 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 import sys from mpl_toolkits import mplot3d sys.path.insert(0, '../../') import fourierOptics as fouOpt class FourierOptics: def __init__(self, l): self.l = l # Wavelength def plane_wave_propagation_rotation(self, A, dpix=None, Rx=None, Ry=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.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) # # Plotting stuff # if False: # Rx0 = Rx # Ry0 = Ry # # Figure 3 from Matsushima et al., 2003 # fig = plt.figure() # ax = fig.add_subplot(121) # ax.grid(alpha=0.3) # ax.set_xlabel('u') # ax.set_ylabel('$\\hat{u}$') # ax.set_aspect('equal') # ax2 = fig.add_subplot(122) # ax2.grid(alpha=0.3) # ax2.set_xlabel('v') # ax2.set_ylabel('$\\hat{v}$') # ax2.set_aspect('equal') # for theta in [0]:#,30,45,60,90]: # Rx = 60 # Ry = 60 # print(Rx, Ry) # Rx = np.deg2rad(Rx) # Ry = np.deg2rad(Ry) # _Rx = np.array([[1,0,0], # [0, np.cos(Rx), -np.sin(Rx)], # [0, np.sin(Rx), np.cos(Rx)]]) # _Ry = np.array([[np.cos(Ry), 0, np.sin(Ry)], # [0,1,0], # [-np.sin(Ry), 0, np.cos(Ry)]]) # _R = np.dot(_Ry,_Rx) # print(_R) # _fz = np.sqrt( (self.l**(-2) - _fx**2 - _fy**2).astype('complex128') ) # _fxr = (_R[0,0]*_fx + _R[1,0]*_fy + _R[2,0]*_fz) # _fyr = (_R[0,1]*_fx + _R[1,1]*_fy + _R[2,1]*_fz) # _fzr = np.sqrt( (self.l**(-2) - _fxr**2 - _fyr**2).astype('complex128') ) # _u0 = _R[2,0]*1/self.l # _v0 = _R[2,1]*1/self.l # # Recenter the new rotated frequencies at 0,0 # _fxr_c = (_fxr - _u0) # _fyr_c = (_fyr - _v0) # _fzr_c = np.sqrt( (self.l**(-2) - _fxr_c**2 - _fyr_c**2).astype('complex128') ) # # _Az02 = np.fft.fftshift(_Az0) # _fx_l2 = np.fft.fftshift(_fx_l) # _fy_l2 = np.fft.fftshift(_fy_l) # _fxr2 = np.fft.fftshift(_fxr) # _fyr2 = np.fft.fftshift(_fyr) # _fxr_c2 = np.fft.fftshift(_fxr_c) # _fyr_c2 = np.fft.fftshift(_fyr_c) # ax.plot(_fx_l2, _fxr2[1000], label='$\\phi_y={:.0f}$'.format(np.degrees(Ry))) # ax2.plot(_fy_l2, _fyr2[:,1000], label='$\\phi_x={:.0f}$'.format(np.degrees(Rx))) # ax.legend() # ax2.legend() # ax.set_xlim([-50,50]) # ax.set_ylim([-50,50]) # ax2.set_xlim([-50,50]) # ax2.set_ylim([-50,50]) # plt.tight_layout() # plt.show() # Rx = Rx0 # Ry = Ry0 # Figure showing frequency shift after Fourier rotatation # fig = plt.figure(figsize=(8,3.5)) # ax = fig.add_subplot(111) # ax.plot(_fyr2[:,1000], np.abs(_Az02)[:,1000], 'k-', label='Reference frame ($\\hat{u}$)') # ax.plot(_fy_l2, np.abs(_Az02)[:,1000], 'k-', alpha=0.4, label='Source frame ($u$)') # ax.set_xlim([-0.5,+0.5]) # ax.set_ylim(bottom=0) # ax.set_xlabel('Frequency') # ax.set_ylabel('Amplitude') # ax.axes.get_xaxis().set_ticks([_v0,0]) # ax.axes.xaxis.set_ticklabels(['$\\hat{u}=u_0$', 0]) # ax.axes.get_yaxis().set_ticks([]) # plt.grid(alpha=0.2) # plt.legend() # plt.savefig('frequency_shift_illustration.pdf', bbox_inches='tight', transparent=True) # plt.show() Rx = np.deg2rad(Rx) Ry = np.deg2rad(Ry) _Rx = np.array([[1,0,0], [0, np.cos(Rx), -np.sin(Rx)], [0, np.sin(Rx), np.cos(Rx)]]) _Ry = np.array([[np.cos(Ry), 0, np.sin(Ry)], [0,1,0], [-np.sin(Ry), 0, np.cos(Ry)]]) _R = np.dot(_Rx,_Ry) _fz = np.sqrt( (self.l**(-2) - _fx**2 - _fy**2).astype('complex128') ) _fxr = (_R[0,0]*_fx + _R[1,0]*_fy + _R[2,0]*_fz) _fyr = (_R[0,1]*_fx + _R[1,1]*_fy + _R[2,1]*_fz) _fzr = np.sqrt( (self.l**(-2) - _fxr**2 - _fyr**2).astype('complex128') ) _u0 = _R[2,0]*1/self.l _v0 = _R[2,1]*1/self.l # Recenter the new rotated frequencies at 0,0 _fxr_c = (_fxr - _u0) _fyr_c = (_fyr - _v0) _fzr_c = np.sqrt( (self.l**(-2) - _fxr_c**2 - _fyr_c**2).astype('complex128') ) # Calculation of the Jacobian # _J = (_R[0,0]*_R[1,1] - _R[0,1]*_R[1,0]) + (_R[0,1]*_R[1,2] - _R[0,2]*_R[1,1])*_fxr_c/_fzr_c + (_R[0,2]*_R[1,0]-_R[0,0]*_R[1,2])*_fyr_c/_fzr_c _J = (_R[0,0]*_R[1,1] - _R[0,1]*_R[1,0]) + (_R[1,0]*_R[2,0] - _R[2,0]*_R[1,1])*_fxr_c/_fzr_c + (_R[0,2]*_R[2,0]-_R[0,0]*_R[2,1])*_fyr_c/_fzr_c _E_rot = pyfftw.empty_aligned((_n,_m), dtype='complex128') _A_in = pyfftw.empty_aligned((_n,_m), dtype='complex128') _A_in[:] = _Az0 * np.abs(_J) _FFTW_Obj2 = pyfftw.FFTW( _A_in, _E_rot, axes=(0,1), direction='FFTW_BACKWARD', threads=6, flags=('FFTW_ESTIMATE',)) _FFTW_Obj2() # Coordinates of E_xy0 _y, _x = np.mgrid[-_n / 2 * dpix : _n / 2 * dpix : 1j*_n, -_m / 2 * dpix : _m / 2 * dpix : 1j * _m] _carr_freq = np.exp(1j * 2 * np.pi * (_u0 * _x + _v0 * _y)) _E_rot = np.fft.fftshift(_E_rot) # use fftshift to move the origin from edge to center _E_rot *= _carr_freq _E_rot /= np.sqrt(_IntensityNormFactor) # Normalization back to initial energy. if Rx or Ry: # A rotation requires interpolation over the output # Performing interpolation to get to the required final plane, done for real and imaginary part. _xs = np.linspace(-_n/2*dpix, _n/2*dpix, _n) _ys = np.linspace(-_m/2*dpix, _m/2*dpix, _m) _zs = np.zeros(_xs.shape) _xr = _R[0,0]*_xs + _R[1,0]*_ys + _R[2,0]*_zs # Rotated original coordinates _yr = _R[0,1]*_xs + _R[1,1]*_ys + _R[2,1]*_zs # Rotated original coordinates # _zr = _R[0,2]*_xs + _R[1,2]*_ys + _R[2,2]*_zs _f_ip_real = interpolate.interp2d(_xs, _ys, _E_rot.real, kind='linear') _f_ip_imag = interpolate.interp2d(_xs, _ys, _E_rot.imag, kind='linear') # _f_ip_real = interpolate.interp2d(_xr, _yr, _E_rot.real, kind='linear') # _f_ip_imag = interpolate.interp2d(_xr, _yr, _E_rot.imag, kind='linear') _xi = np.linspace(_xr.min(), _xr.max(), _n) _yi = np.linspace(_yr.min(), _yr.max(), _m) # _xi = _xs # _yi = _ys _E_out_ip_real = _f_ip_real(_xi, _yi) _E_out_ip_imag = _f_ip_imag(_xi, _yi) _E_rot = _E_out_ip_real + _E_out_ip_imag*1j _y, _x = np.meshgrid(_yi, _xi) if show or save: _E_rot_extent = [np.min(_x), np.max(_x), np.min(_x), np.max(_x)] fig = plt.figure(figsize=(8,3)) ax = fig.add_subplot(121) im = ax.imshow(np.abs(_E_rot), cmap='hot', extent=_E_rot_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_rot), cmap='hot', extent=_E_rot_extent) ax2.set_xlabel('x [mm]') ax2.set_ylabel('y [mm]') ax.set_xlim([-0.4,0.4]) ax.set_ylim([-0.4,0.4]) ax2.set_xlim([-0.4,0.4]) ax2.set_ylim([-0.4,0.4]) fig.colorbar(im2, ax=ax2) plt.tight_layout() if save: plt.savefig(save, bbox_inches='tight', transparent=True) if show: plt.show() elif not show: plt.close() return np.array([_E_rot, _x, _y]) 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 coordinate_rotation(self, Rx, Ry, n=5, m=5): _n = n _m = m Rx = np.deg2rad(Rx) Ry = np.deg2rad(Ry) dpix = 1/30 * self.l * 1000 _Rx = np.array([[1,0,0], [0, np.cos(Rx), -np.sin(Rx)], [0, np.sin(Rx), np.cos(Rx)]]) _Ry = np.array([[np.cos(Ry), 0, np.sin(Ry)], [0,1,0], [-np.sin(Ry), 0, np.cos(Ry)]]) _R = np.dot(_Rx,_Ry) theta = range(360) _xcirc = 1.2 * np.cos(np.deg2rad(theta)) _ycirc = 1.2 * np.sin(np.deg2rad(theta)) _zcirc = np.zeros(_xcirc.shape) _ys, _xs = np.mgrid[-_n / 2 * dpix : _n / 2 * dpix : 1j*_n, -_m / 2 * dpix : _m / 2 * dpix : 1j * _m] _zs = np.zeros(_xs.shape) _xr = _R[0,0]*_xs + _R[1,0]*_ys + _R[2,0]*_zs # Rotated original coordinates _yr = _R[0,1]*_xs + _R[1,1]*_ys + _R[2,1]*_zs # Rotated original coordinates _zr = _R[0,2]*_xs + _R[1,2]*_ys + _R[2,2]*_zs _xcirc_r = _R[0,0]*_xcirc + _R[1,0]*_ycirc + _R[2,0]*_zcirc # Rotated original coordinates _ycirc_r = _R[0,1]*_xcirc + _R[1,1]*_ycirc + _R[2,1]*_zcirc # Rotated original coordinates _zcirc_r = _R[0,2]*_xcirc + _R[1,2]*_ycirc + _R[2,2]*_zcirc _xcirc_p = _xcirc # Projected original coordinates _ycirc_p = _ycirc # Projected original coordinates _zcirc_p = (-_R[2,0]*_xcirc_p - _R[2,1]*_ycirc_p) / _R[2,2] # Plot 3d Wireframe of source plane and reference plane fig = plt.figure(figsize=(8,8)) ax = plt.axes(projection='3d') ax.plot_wireframe(_zs, _xs, _ys, label='source frame', color='blue', alpha=.3) ax.plot_wireframe(_zr, _xr, _yr, label='reference frame', color='red', alpha=.3) ax.plot3D(_zcirc, _xcirc, _ycirc, 'blue') ax.plot3D(_zcirc_r, _xcirc_r, _ycirc_r, 'red', alpha=.2) ax.plot3D(_zcirc_p, _xcirc_p, _ycirc_p, 'black') ax.set_xlabel('z') ax.set_ylabel('x') ax.set_zlabel('y') ax.set_xlim([-3,3]) ax.set_ylim([-3,3]) ax.set_zlim([-3,3]) ax.grid(True, alpha=0.3) ax.legend() plt.tight_layout() # plt.savefig('./coordinate_rotation_figures/coordinate_rotation_rx{:.0f}_ry{:.0f}.png'.format(np.degrees(Rx), np.degrees(Ry)), bbox_inches='tight') # plt.close() plt.show(block=False) # _Az02 = np.fft.fftshift(_Az0) # _fx_l2 = np.fft.fftshift(_fx_l) # _fy_l2 = np.fft.fftshift(_fy_l) # _fxr2 = np.fft.fftshift(_fxr) # _fyr2 = np.fft.fftshift(_fyr) # plt.plot(_fyr2[:,1000], np.abs(_Az02)[:,1000]); plt.plot(_fy_l2, np.abs(_Az02)[:,1000]); plt.grid(); plt.show() 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) fouOpt_norm = fouOpt.FourierOptics(l=l) # Pick sampling parameters dpix = 1/40*l n, m = 3001, 3001 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 = fouOpt.Profiles() # A = prof.spiderTopHat(250, 100, 2001, 2001, dpix=1, oversample=1) # A = prof.top_hat(4*l, n, m, dpix=dpix, oversample=1) x, y = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m)) gauss = fouOpt.GaussianBeam(0.02, l) A = gauss.fundamental_mode(0,n,m,np.max(x)) # z0 = n*dpix * np.tan(np.deg2rad(30)) # A = gauss.fundamental_mode(z0,n,m,np.max(x)) # Exyz = A # # Calculate the Gaussian beam for the whole tilted image plane # A_at_z4l = gauss.fundamental_mode(4*l,n,m,np.max(x)) # The input distribution gets a curved wavefront after the lens U_in = A # sys.exit() Exyz,x4l,y4l = fouOpt_norm.plane_wave_propagation(A, dz=4*l, dpix=dpix, oversample=None, show=True) fou.coordinate_rotation(45,20) # _E_res_abs = np.abs(A_at_z4l)/np.max(np.abs(A_at_z4l)) - np.abs(Exyz)/np.max(np.abs(Exyz)) # _E_res_ang = np.angle(A_at_z4l)/np.max(np.angle(A_at_z4l)) - np.angle(Exyz)/np.max(np.angle(Exyz)) # _E_rot_extent = [np.min(x4l), np.max(x4l), np.min(y4l), np.max(x4l)] # fig = plt.figure(figsize=(8,3)) # ax = fig.add_subplot(121) # im = ax.imshow(_E_res_abs.T, cmap='hot', extent=_E_rot_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(_E_res_ang.T, cmap='hot', extent=_E_rot_extent) # ax2.set_xlabel('x [mm]') # ax2.set_ylabel('y [mm]') # ax.set_xlim([-0.4,0.4]) # ax.set_ylim([-0.4,0.4]) # ax2.set_xlim([-0.4,0.4]) # ax2.set_ylim([-0.4,0.4]) # fig.colorbar(im2, ax=ax2) # plt.tight_layout() # # plt.show()/ # plt.close() Erot2,x,y = fou.plane_wave_propagation_rotation(Exyz, dpix=dpix, Rx=45, Ry=20, oversample=None, show=True)#, save='p2p_gauss_Rx45_Ry45.png') ''' # Compare phase to gaussian analysis. z0 = n*dpix * np.tan(np.deg2rad(30)) z = np.linspace(0, 2*z0, n) print(int((3+np.sqrt(3))/3*n)) datacube = [] for i in range(len(z)): print(i) A = gauss.fundamental_mode(z[i],n,m,np.max(x)) # _f_ip_real = interpolate.interp2d(x, y, A.real, kind='linear') # _f_ip_imag = interpolate.interp2d(x, y, A.imag, kind='linear') # A_out_ip_real = _f_ip_real(_xi, _yi) # A_out_ip_imag = _f_ip_imag(_xi, _yi) # A = _E_out_ip_real + _E_out_ip_imag*1j datacube.append(A[:,i]) plt.imshow(np.angle(np.asarray(datacube)));plt.colorbar();plt.show() plt.imshow(np.abs(np.asarray(datacube)));plt.colorbar();plt.show() set_trace() '''