Files
fourierOptics/scripts/planewave_rotation_test.py
T

467 lines
16 KiB
Python

#!/usr/bin/env Python3
import numpy as np
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, '../')
sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Documents/Published/FourierOptics GitHub release/')
import fourierOptics as fouOpt
'''
Description: Old script when I was trying to make propagation to off-axis parabolic mirrors
work in 3D. This was based on papers by Matsushima (around 2004, I believe).
The script is currently broken.
If you get this to work, please do let me know!
Contact: born@astron.nl
'''
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') )
set_trace()
# 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.
_xls = np.linspace(-_n/2*dpix, _n/2*dpix, _n)
_yls = np.linspace(-_m/2*dpix, _m/2*dpix, _m)
_zls = np.zeros(_xls.shape)
_x = _R[0,0]*_xls + _R[1,0]*_yls + _R[2,0]*_zls # Rotated original coordinates
_y = _R[0,1]*_xls + _R[1,1]*_yls + _R[2,1]*_zls # Rotated original coordinates
_z = _R[0,2]*_xls + _R[1,2]*_yls + _R[2,2]*_zls
_ys, _xs = np.mgrid[-_n / 2 * dpix : _n / 2 * dpix : 1j*_n, -_m / 2 * dpix : _m / 2 * dpix : 1j * _m]
# _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
# fc = np.fft.fftshift(np.angle(_E_xy0))[::3,::3]
# fcr = np.angle(_E_rot)[::3,::3]
# norm=colors.Normalize()
# fig = plt.figure(figsize=(8,8))
# ax = plt.axes(projection='3d')
# # ax.plot_surface(_zs[::3,::3], _xs[::3,::3], _ys[::3,::3], facecolors=cm.hot(norm(fc)), alpha=0.5)
# ax.plot_surface(_zr[::3,::3], _xr[::3,::3], _yr[::3,::3], facecolors=cm.hot(norm(fcr)))
# ax.set_xlabel('z')
# ax.set_ylabel('x')
# ax.set_zlabel('y')
# ax.set_xlim([-2,2])
# ax.set_ylim([-2,2])
# ax.set_zlim([-2,2])
# plt.show()
_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
# set_trace()
_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)
# set_trace()
if show or save:
_E_rot_extent = [np.min(_x), np.max(_x), np.min(_y), 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[0,2]*_xcirc*np.sqrt(1+np.tan(Ry)**2) + _R[1,2]*_ycirc*np.sqrt(1+np.tan(Rx)**2)
_zcirc_p = np.tan(Ry)*_xcirc - np.tan(Rx)*_ycirc + _zcirc
# _zcirc_p = -_xcirc / np.tan(Ry) - _ycirc / np.tan(Rx) + _zcirc
# 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=.1)
ax.plot_wireframe(_zr, _xr, _yr, label='reference frame', color='red', alpha=.1)
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()
# _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/30*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(.4*l,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
# fou.coordinate_rotation(45,30)
# sys.exit()
Exyz,x4l,y4l = fouOpt_norm.plane_wave_propagation(A, dz=4*l, dpix=dpix, oversample=None, show=False)
# _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()
# Erot,_,_ = fou.plane_wave_propagation_rotation(Exyz, dpix=dpix, Rx=0, Ry=45, oversample=None, show=False)
# Erot2,_,_ = fou.plane_wave_propagation_rotation(Exyz, dpix=dpix, Rx=0, Ry=45, oversample=None, show=True)
Erot2,x,y = fou.plane_wave_propagation_rotation(Exyz, dpix=dpix, Rx=60, Ry=0, oversample=None, show=True)#, save='p2p_gauss_Rx45_Ry45.png')