mirror of
https://github.com/yosata/fourierOptics.git
synced 2026-10-29 20:51:10 +00:00
362 lines
12 KiB
Python
362 lines
12 KiB
Python
#!/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()
|
|
'''
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|