#!/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, '../../') # where fourieroptics.py is located, so that it can be imported 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) _fz = np.sqrt( (self.l**(-2) - _fx**2 - _fy**2).astype('complex128') ) # Rotation about the x and y axis from radians to degrees Rx = np.deg2rad(Rx) Ry = np.deg2rad(Ry) # Rotation matrices _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) # Rotation in k-space _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') ) # Carrier frequency _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 # Fourier transform of the input distribution, _A_in _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_rot _y, _x = np.mgrid[-_n / 2 * dpix : _n / 2 * dpix : 1j*_n, -_m / 2 * dpix : _m / 2 * dpix : 1j * _m] # Multiplication with the phase contribution due to frequency shift, # i.e. shift theorem. _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 _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: # Plot the amplitude and phase _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): # Visualisation of the rotation matrix and projection. # # A circle C is plotted on the source frame (dark blue), # the rotation of C onto the reference plane is shown in red # and C is also projected from the source frame onto the # reference frame in black. # _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) # A grid in the source frame, coordinates (xs, ys, zs) _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) # Rotated coordinates of the grid in the reference frame, coordinates (xr, yr, zr) _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 # The circle C in the source frame. theta = range(360) _xcirc = 1.2 * np.cos(np.deg2rad(theta)) _ycirc = 1.2 * np.sin(np.deg2rad(theta)) _zcirc = np.zeros(_xcirc.shape) # Rotated coordinates _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 # Projected coordinates _xcirc_p = _xcirc _ycirc_p = _ycirc _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) if __name__ == '__main__': # Pick a wavelength l = 30e-3 # Initiate Fourier Optics fou = FourierOptics(l=l) fouOpt_norm = fouOpt.FourierOptics(l=l) # Pick sampling parameters dpix = 1/15*l n, m = 1001, 1001 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() # Use a Top Hat profile: # A = prof.top_hat(4*l, n, m, dpix=dpix, oversample=1) # Or use a gaussian: 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.06, l) A = gauss.fundamental_mode(4*l,n,m,np.max(x)) # Definition of the rotation angles. # Rotation about the x-axis is done first, then about the y-axis. theta_x = 45 # Degrees theta_y = 45 # Degrees # Plane wave propagation to 4*lambda # Exyz,x4l,y4l = fouOpt_norm.plane_wave_propagation(A, dz=4*l, dpix=dpix, show=True) # Show the input distribution _E_rot_extent = [np.min(x), np.max(x), np.min(y), np.max(y)] fig = plt.figure(figsize=(8,3)) ax = fig.add_subplot(121) im = ax.imshow(np.abs(A), cmap='hot', extent=_E_rot_extent) 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(A), cmap='hot', extent=_E_rot_extent) ax2.set_xlabel('x [mm]') ax2.set_ylabel('y [mm]') fig.colorbar(im2, ax=ax2) plt.tight_layout() plt.show() # Illustration of the rotation and projection for the given rotation angles. fou.coordinate_rotation(theta_x,theta_y) # Plane wave propagation onto a rotated plane. Erot,x,y = fou.plane_wave_propagation_rotation(A, dpix=dpix, Rx=theta_x, Ry=theta_y, show=True) ''' # 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() '''