yosata/fourierOptics
1
0 Fork 0
mirror of https://github.com/yosata/fourierOptics.git synced 2026-10-29 20:51:10 +00:00
Code Issues Projects Releases Wiki Activity

Initial commit

Browse Source
This commit is contained in:
yosata committed 2020-05-28 09:41:08 +02:00
1 parent 91aeef5a90
commit f4cc304953
23 files changed
+8948 -27
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.gitignore
+3 -27
View File
@@ -1,30 +1,6 @@
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
pip-wheel-metadata/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
.DS_Store
__pychache__/
.ipynb_checkpoints/
MANIFEST
# PyInstaller
fourierOptics/__init__.py
+9
View File
@@ -0,0 +1,9 @@
import os
from .fourier import *
from .gaussian import *
from .opticalTrain import *
from .optics import *
from .profiles import *
from .psfFitter import *
from .wavefrontErrors import *
from .misc import *
fourierOptics/data/ELT_M1_pupil_profile.fits
Executable
BIN
View File
Binary file not shown.
fourierOptics/data/__init__.py
View File
Whitespace-only changes.
fourierOptics/data/__pycache__/__init__.cpython-37.pyc
BIN
View File
Binary file not shown.
fourierOptics/data/__pycache__/glassCatalog.cpython-37.pyc
BIN
View File
Binary file not shown.
fourierOptics/data/glassCatalog.py
+12
View File
@@ -0,0 +1,12 @@
#!/usr/bin/env Python3
# Pyton dictionary with some glass types
catalog = {
'N-BK7' : [1.039612120E+00, 6.000698670E-03, 2.317923440E-01, 2.001791440E-02, 1.010469450E+00, 1.035606530E+02 ],
'SiO2' : [0.6961663, 0.0684043, 0.4079426, 0.1162414, 0.8974794, 9.896161 ],
'F2' : [1.345333590E+00, 9.977438710E-03, 2.090731760E-01, 4.704507670E-02, 9.373571620E-01, 1.118867640E+02 ],
'S-FPL51' : [1.170105050E+00, 6.162039240E-03, 4.757107830E-02, 2.633728760E-02, 7.638324450E-01, 1.418826420E+02 ],
'S-LAH71' : [1.982800310E+00, 1.189874560E-02, 3.167584500E-01, 5.271560010E-02, 2.444726460E+00, 2.132206970E+02 ],
}
fourierOptics/fourier.py
+440
View File
@@ -0,0 +1,440 @@
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
class FourierOptics:
def __init__(self, l):
self.l = l # Wavelength
def plane_wave_propagation(self, A, xin, yin, dz=None, oversample=None, show=False, save=False, scaleInputArray=False, outputCoordinates=None):
'''
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)
# Calculate the pixel size from the xi, eta plane
_dpix = np.mean(np.diff(xin))
if _dpix < 1e-99:
_dpix = np.mean(np.diff(xin))
# 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)
_kz = self.phase_exponent(_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]
_y, _x = yin, xin
if outputCoordinates is not None:
# 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(outputCoordinates[0],outputCoordinates[1],_n)
_yi = np.linspace(outputCoordinates[2],outputCoordinates[3],_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[outputCoordinates[0]:outputCoordinates[1]:1j*_n, outputCoordinates[2]:outputCoordinates[3]:1j * _m]
_E_xyz_extent = [np.min(_x.real), np.max(_x.real), np.min(_y.real), np.max(_x.real)]
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.real, _y.real])
def fresnel_propagation(self, A, xin, yin, dz=None, oversample=None, show=False, save=False,
scaleInputArray=False, outputCoordinates=None, apply_interpolation=True):
'''
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.beam_radius(dz)[0]
_xip = np.linspace(-_halfwidth, _halfwidth, _n) # Desired x,y - plane
_yip = np.linspace(-_halfwidth, _halfwidth, _m) # Desired x,y - plane
elif outputCoordinates is not None:
_xip = np.linspace(outputCoordinates[0], outputCoordinates[1], _n) # Desired x,y - plane
_yip = np.linspace(outputCoordinates[2], outputCoordinates[3], _m) # Desired x,y - plane
else:
_xip = xin[0]
_yip = yin[:,0]
# 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.
if outputCoordinates and apply_interpolation:
_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(((_x.max()/_xip.max())*(_y.max()/_yip.max())))
_E_out_ip = np.conj(_E_out_ip)
elif outputCoordinates and not apply_interpolation:
_arraySelection = self.select_subsection(_x, _y, outputCoordinates)
_xip = _x.real[_arraySelection[0]:_arraySelection[1]+1]
_yip = _y.real[_arraySelection[2]:_arraySelection[3]+1]
_E_out_ip = _E_out[_arraySelection[0]:_arraySelection[1]+1,_arraySelection[2]:_arraySelection[3]+1]
_E_out_ip = np.conj(_E_out_ip)
else:
_E_out_ip = _E_out
_xip, _yip = _x, _y
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.real, _y.real])
def freq_array(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 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 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.gaussian_beam_fit(a,dz,x,y) - og_data).flatten()).real
def gaussian_beam_fit(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
def optical_transfer_function(self, PSF, x, y, dpix=None):
'''
Calculates the optical transfer function of the system response
Input is the image on the detector, preferably of a PSF
The frequencies are given over the whole provided range, but are
typically limited to 1 / (lambda * F/#).
'''
if dpix is None:
dpix = np.mean(np.diff(x.real))
if dpix < 1e-99:
dpix = np.mean(np.diff((x.real).T))
print('dpix is {}'.format(dpix))
# FFT input
_A_in = pyfftw.empty_aligned((PSF.shape), dtype='complex128')
_A_in[:] = PSF
# Output of the FFT
_A_out = pyfftw.empty_aligned(PSF.shape, dtype='complex128')
_FFTW_Obj = pyfftw.FFTW( np.fft.ifftshift(_A_in), _A_out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
_FFTW_Obj()
_A_out = np.fft.fftshift(_A_out)
# Normalise the output
_OTF = _A_out / np.max(_A_out)
_fx = np.fft.fftshift(np.fft.fftfreq(np.shape(x)[0], d=dpix))
_fy = np.fft.fftshift(np.fft.fftfreq(np.shape(y)[0], d=dpix))
return _OTF, _fx, _fy
def modulation_transfer_function(self, *args):
'''
Calculates the modulation transfer function of the system response
Input is the image on the detector, preferably of a PSF
The frequencies are given over the whole provided range, but are
typically limited to 1 / (lambda * F/#).
'''
_OTF, _fx, _fy = self.optical_transfer_function(*args)
return np.abs(_OTF), _fx, _fy
def encircled_energy_function(self, PSF, x, y, x0=0, y0=0, dpix=None, originalPSF=None):
'''
Function that returns the encircled energy as a function of distance
from the centroid position, defined by x0 and y0.
Parameters:
PSF - Intensity map of the PSF
x,y - x and y coordinates corresponding to the PSF
x0,y0 - Optional; location of the centroid. If not given,
x0=0 and y0=0 are assumed.
dpix - Optional; pixel scale, default=1
Returns:
_EEF - Array with encircled energy values
_r - Array with radius values
'''
if dpix is None:
dpix = np.mean(np.diff(x))
if dpix < 1e-99:
dpix = np.mean(np.diff(x.T))
if originalPSF is None:
_total_energy = np.sum(PSF)
else:
_total_energy = np.sum(originalPSF)
_r = np.sqrt((x.real-x0)**2 + (y.real-y0)**2)
_r_ls = np.linspace(0,_r.max(),len(_r))
_mask = [ _r <= i for i in _r_ls]
_EEF = np.asarray([np.sum(PSF[m]) for m in _mask])/_total_energy
_r = _r_ls * dpix
return _EEF, _r
@staticmethod
def find_nearest(array, value):
array = np.asarray(array)
diff = (np.abs(array - value)).argmin()
idx = np.unravel_index(diff, array.shape)
return idx
def select_subsection(self, x, y , outputCoordinates):
select_1 = np.max(self.find_nearest(x, outputCoordinates[0]))
select_2 = np.max(self.find_nearest(x, outputCoordinates[1]))
select_3 = np.max(self.find_nearest(y, outputCoordinates[2]))
select_4 = np.max(self.find_nearest(y, outputCoordinates[3]))
return [select_1, select_2, select_3, select_4]
fourierOptics/fourierOptics.py
+2095
View File
File diff suppressed because it is too large. Load diff
fourierOptics/gaussian.py
+113
View File
@@ -0,0 +1,113 @@
import numpy as np
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 fundamental_mode(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.beam_radius(z)
_R = self.curvature(z)
_phi0 = self.phase_shift(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 beam_radius(self, z):
return self.w0 * np.sqrt(1 + (z/self.z_c)**2)
def phase_shift(self, z):
return np.arctan(z/self.z_c)
def relative_power(self, r, z):
return np.exp(-2 * (r / self.beam_radius(z))**2)
class GaussianBeamPropagation(GaussianBeam):
def __init__(self, gaussianBeamObject):
self.inputBeam = gaussianBeamObject
self.l = self.inputBeam.l
self.z_c = self.inputBeam.z_c
self.theta_0 = self.inputBeam.theta_0
self.w0 = self.inputBeam.w0
self.fwhm = self.inputBeam.fwhm
self.beam_width = self.w0
def propagate_with_a_lens(self, ABCD, d_out, n, m, d_in=0, rmax=None):
'''
Propagates the gaussian beam through a lens with the ray transfer matrix ABCD.
This follows the equations of Ch. 3.3 of Goldsmith's Quasi-optical Systems
Parameters:
ABCD - Ray transfer matrix, use ABCD_paraxial_lens to calculate
this matrix for a paraxial lens, with focal length f.
d_out - Distance after the lens over which the beam is propagated.
d_in - Distance from the beam waist (w0) to the lens. Default is 0.
rmax - Maximum radius for which to calculate the beam profile.
If None, it will be set to three times the beam width at
the output distance.
Output:
The beam information. Can be plotted using imshow()
The beam width at the given point can be retrieved by using self.beam_width
'''
_A,_B,_C,_D = ABCD
_focalDistance = - ( (_A * d_in + _B) * (_C * d_in + _D) + _A * _C * self.z_c**2 ) / ( (_C * d_in + _D)**2 + _C**2 * self.z_c**2 )
_z = d_out - _focalDistance
if d_out > _focalDistance:
d_l = d_out - _focalDistance
d_out = _focalDistance
_w0_out = self.w0 / np.sqrt((_C * d_in + _D)**2 + _C**2 * self.z_c**2)
_z_c_new = np.pi * _w0_out**2 / self.l
_wl_out = _w0_out * np.sqrt(1 + (d_l/_z_c_new)**2)
# _q_out = ((_A + _C * d_out )*1j*self.z_c + (_A + _C * d_out)*d_in + (_B + _D*d_out)) / (_C * 1j * self.z_c + _C * d_in + _D)
_w0_out = self.w0 / np.sqrt((_C * d_in + _D)**2 + _C**2 * self.z_c**2)
self.z_c = np.pi * _w0_out**2 / self.l
self.w0 = _w0_out
self.theta_0 = self.l / (np.pi * _w0_out)
self.fwhm = 1.18 * self.theta_0
self.beam_width = self.beam_radius(_z)
if rmax is None:
rmax = 3 * self.beam_width
_y, _x = np.mgrid[-rmax:rmax:n*1j,-rmax:rmax:n*1j]
_r = np.sqrt(_x**2 + _y**2)
if _z != 0:
_w = self.beam_radius(_z)
_R = self.curvature(_z)
_phi0 = self.phase_shift(_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 ABCD_paraxial_lens(self, f):
return [1, 0, -1/f, 1]
fourierOptics/misc.py
+95
View File
@@ -0,0 +1,95 @@
import numpy as np
import pickle
'''
Miscellaneous functions that do not belong specifically to any of the
other files.
Author: Joost van den Born
Contact: born@astron.nl
'''
def perform_zero_padding(A, xs, ys):
'''
Add zero padding around the array
Input:
A - 2D-Input array
xs,ys - The new shape after padding.
Output:
- Array with the new padded zeros
'''
assert A.shape[0] <= xs and A.shape[1] <= ys, 'New size is smaller than array size.'
x_margin = int((xs - A.shape[0])/2)
y_margin = int((ys - A.shape[1])/2)
return np.pad(A, (x_margin, y_margin))
def find_nearest(array, value):
'''
For a value it finds the nearest element in an array
Returns the index of that element in the original
array.
'''
array = np.asarray(array)
diff = (np.abs(array - value)).argmin()
idx = np.unravel_index(diff, array.shape)
return idx
def find_maximum(array):
'''
Finds the index of the maximum of the given array
'''
idx = np.unravel_index(array.argmax(), array.shape)
return idx
def select_subsection(x, y , outputCoordinates):
'''
Returns the indices of a subsection of an x,y meshgrid
'''
select_1 = find_nearest(x, outputCoordinates[0])[1]
select_2 = find_nearest(x, outputCoordinates[1])[1]
select_3 = find_nearest(y, outputCoordinates[2])[0]
select_4 = find_nearest(y, outputCoordinates[3])[0]
return [select_1, select_2, select_3, select_4]
def perform_extrapolation(x, y, n, m):
'''
Performs linear extrapolation of a meshgrid to a new
shape n,m
Returns the new extrapolated x,y meshgrid
'''
# Determine properties of x,y meshgrid
xshape = x.shape[0]
yshape = y.shape[0]
dpix = (np.abs(np.min(x)) + np.abs(np.max(x))) / (xshape-1)
# Extrapolate from, and include, the old x,y meshgrid
xnew = np.linspace(-(n-1)/2*dpix, (n-1)/2*dpix, n)
ynew = np.linspace(-(m-1)/2*dpix, (m-1)/2*dpix, m)
xnew, ynew = np.meshgrid(xnew, ynew)
return xnew, ynew
def save_variable(filename, variable):
'''
Dumps the given variable to a binary pickle file
for later use.
'''
with open(filename, 'wb') as f:
pickle.dump(variable, f, protocol=-1 )
def load_variable(filename):
'''
Loads a binary pickle file for use.
'''
with open(filename, 'rb') as f:
variable = pickle.load(f)
return variable
fourierOptics/opticalTrain.py
+103
View File
@@ -0,0 +1,103 @@
import numpy as np
from tabulate import tabulate
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.fresnel_propagation(_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.fresnel_propagation(_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])
fourierOptics/optics.py
+773
View File
@@ -0,0 +1,773 @@
import numpy as np
from scipy.optimize import fsolve
from .data import glassCatalog
class LensObject:
def __init__(self, CA, name=None):
if name:
assert type(name)==str, 'No string for a name.'
self.name = name
else:
self.name = str(self)
self.CA = CA
self.type = None
self.f = None
self.sag = None
self.wavefrontError = False
def add_wavefront_error(self, z):
assert len(z.shape)==2, 'Wavefront data is not 2D. Cannot continue'
self.sag = z
self.wavefrontError = True
class ParaxialLens(LensObject):
def __init__(self, f, CA, name=None):
super().__init__(CA, name)
self.f = f
self.type = 'PARAXIAL'
def find_focal_length(self):
return self.f
def phase_delay(self, x, y, l, **kwargs):
_k = 2 * np.pi / l
_t_l = np.exp(1j * _k / (2 * self.f) * (x**2 + y**2))
if self.wavefrontError:
_t_l *= np.exp(-1j * _k * self.sag / l)
# 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/2)).astype(int)
# _t_l *= _lensMask
return _t_l
class SphericalLens(LensObject):
'''
The SphericalLens object describes any spherical lens and
finds the phase delay for all light traveling through it.
The refractive index, of course, depends on the wavelength.
'''
def __init__(self, glass, thickness, R1, R2, CA, name=None):
super().__init__(CA, name)
assert glassCatalog.catalog.get(glass), 'Glass type not in catalog'
self.type = 'SPHERICAL'
self.thickness = thickness # mm
self.R1 = R1 # mm
self.R2 = R2 # mm
self.glass = glass
self.dispersionCoefficients = glassCatalog.catalog.get(self.glass)
self.referenceWavelength = 1.35e-3 # mm
self.f = self.find_focal_length()
def set_first_surface_sag(self, x,y):
'''
Defines the shape of the first surface and bound the corresponding sag function to
the lens object.
Input:
x,y - Coordinates of the sag.
'''
# A spherical lens surface requires only the radius of curvature
_first_surface_sag = self.R1 * (1 - np.sqrt(1 - (x**2 + y**2) / self.R1**2 ) )
dist_from_center = np.sqrt(x**2 + y**2)
_lensMask = (dist_from_center <= (self.CA/2)).astype(int)
_first_surface_sag *= _lensMask
self.first_surface_sag = _first_surface_sag
return self.first_surface_sag
def set_second_surface_sag(self, x,y):
'''
Defines the shape of the second surface and bound the corresponding sag function to
the lens object.
Input:
x,y - Coordinates of the sag.
'''
# A spherical lens surface requires only the radius of curvature
_second_surface_sag = self.R2 * (1 - np.sqrt(1 - (x**2 + y**2) / self.R2**2 ) )
dist_from_center = np.sqrt(x**2 + y**2)
_lensMask = (dist_from_center <= (self.CA/2)).astype(int)
_second_surface_sag *= _lensMask
self.second_surface_sag = _second_surface_sag
return self.second_surface_sag
def get_first_surface_sag(self):
return self.first_surface_sag
def get_second_surface_sag(self):
return self.second_surface_sag
def find_refractive_index(self, l):
# l in mm needs to be converted to l in micron
l *= 1e3
_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 find_focal_length(self):
_n = self.find_refractive_index(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):
'''
This function calculates the phase delay introduced by the lens.
Input:
x, y - Coordinates of the lens
l - Wavelength in millimeters
paraxial - [bool] Apply a paraxial approximation?
'''
_k = 2 * np.pi / l
_n = self.find_refractive_index(l)
if not paraxial:
# For all non-paraxial lenses:
# Retrieve the sag of the lens surfaces for coordinates x,y
_A1 = self.set_first_surface_sag(x,y)
_A2 = self.set_second_surface_sag(x,y)
_dxy = self.thickness - _A1 + _A2
_t_l = np.exp(-1j * _k * self.thickness) * np.exp(-1j * _k * (_n - 1) * _dxy)
elif paraxial:
_t_l = np.exp(-1j * _k * self.thickness) * np.exp(1j * _k / (2 * self.f) * (x**2 + y**2))
if self.wavefrontError:
_t_l *= np.exp(-1j * _k * (_n - 1) * self.sag)
# 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/2)).astype(int)
_t_l *= _lensMask
return _t_l
class EvenAsphereLens(LensObject):
'''
The EvenAsphereLens object describes an even aspherical lens and
finds the phase delay for all light traveling through it.
The refractive index, of course, depends on the wavelength.
'''
def __init__(self, glass, thickness, R1, R2, k1, k2, a1, a2, CA, name=None):
super().__init__(CA, name)
assert glassCatalog.catalog.get(glass), 'Glass type not in catalog'
self.type = 'EVEN ASPHERE'
self.thickness = thickness # mm
self.R1 = R1 # mm
self.R2 = R2 # mm
self.conicConstantS1 = k1
self.conicConstantS2 = k2
self.asphereCoefficientS1 = a1
self.asphereCoefficientS2 = a2
self.glass = glass
self.dispersionCoefficients = glassCatalog.catalog.get(self.glass)
self.referenceWavelength = 1.35 # micron
self.f = self.find_focal_length()
@staticmethod
def calculate_surface_sag(x, y, R, k, a):
'''
Defines the sag shape of the lens
Follows the ZEMAX definition on Even Aspheres
Input parameters:
x, y - x,y coordinates of the lens in mm
R - Radius of curvature of the surface in mm
k - Conic constant of the surface
a - Coefficients of the polynomial expansion of the lens sag.
'''
_r = np.sqrt(x**2 + y**2)
_c = 1 / R
_expansion = np.sum(np.asarray([a[i] * _r**( 2*(i+1) ) for i in range(len(a))]))
_z = _c * _r**2 / ( 1 + np.sqrt( 1 - (1 + k) * _c**2 * _r**2) ) + _expansion
return _z
def set_first_surface_sag(self, x,y):
'''
Defines the shape of the first surface and bound the corresponding sag function to
the lens object.
Input:
x,y - Coordinates of the sag.
'''
# A spherical lens surface requires only the radius of curvature
_first_surface_sag = self.calculate_surface_sag(x, y, self.R1, self.conicConstantS1, self.asphereCoefficientS1)
dist_from_center = np.sqrt(x**2 + y**2)
_lensMask = (dist_from_center <= (self.CA/2)).astype(int)
_first_surface_sag *= _lensMask
self.first_surface_sag = _first_surface_sag
return self.first_surface_sag
def set_second_surface_sag(self, x,y):
'''
Defines the shape of the second surface and bound the corresponding sag function to
the lens object.
Input:
x,y - Coordinates of the sag.
'''
# A spherical lens surface requires only the radius of curvature
_second_surface_sag = self.calculate_surface_sag(x, y, self.R2, self.conicConstantS2, self.asphereCoefficientS2)
dist_from_center = np.sqrt(x**2 + y**2)
_lensMask = (dist_from_center <= (self.CA/2)).astype(int)
_second_surface_sag *= _lensMask
self.second_surface_sag = _second_surface_sag
return self.second_surface_sag
def get_first_surface_sag(self):
return self.first_surface_sag
def get_second_surface_sag(self):
return self.second_surface_sag
def find_refractive_index(self, l):
# l in mm needs to be converted to l in micron
l *= 1e3
_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 find_focal_length(self):
'''
Calculates the focal length, where we assume that the
asphere component is small, so that the focal length is
approximately the same as for the spherical lens with
radii of curvature of R1 and R2.
'''
_n = self.find_refractive_index(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):
'''
This function calculates the phase delay introduced by the lens.
Input:
x, y - Coordinates of the lens
l - Wavelength in millimeters
paraxial - [bool] Apply a paraxial approximation?
'''
_k = 2 * np.pi / l
_n = self.find_refractive_index(l)
if not paraxial:
# For all non-paraxial lenses:
# Retrieve the sag of the lens surfaces for coordinates x,y
_A1 = self.set_first_surface_sag(x,y)
_A2 = self.set_second_surface_sag(x,y)
_dxy = self.thickness - _A1 + _A2
_t_l = np.exp(-1j * _k * self.thickness) * np.exp(-1j * _k * (_n - 1) * _dxy)
elif paraxial:
_t_l = np.exp(-1j * _k * self.thickness) * np.exp(1j * _k / (2 * self.f) * (x**2 + y**2))
if self.wavefrontError:
_t_l *= np.exp(-1j * _k * (_n - 1) * self.sag)
# 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/2)).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)
self.type = None
self.sag = None
self.wavefrontError = False
def add_wavefront_error(self, z):
assert len(z.shape)==2, 'Wavefront data does not have the proper shape. Cannot continue'
self.sag = z
self.wavefrontError = True
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, **kwargs):
_k = 2 * np.pi / l
_dz = (x**2 + y**2) / (4 * self.f)
if self.wavefrontError:
_sagPhase = self.sag
_t_m = np.exp(1j * _k * 2 * (_dz + _sagPhase) )
else:
_t_m = np.exp(1j * _k * 2 * _dz)
# Masking out the parts that fall outside the mirror clear aperture
dist_from_center = np.sqrt(x**2 + y**2)
_mirrorMask = (dist_from_center <= (self.CA/2)).astype(int)
_t_m *= _mirrorMask
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.f = R/2
self.CA = CA
self.type = 'SPHERICAL MIRROR'
def phase_delay(self, x, y, l, **kwargs):
_k = 2 * np.pi / l
_dz = self.R * (1 - np.sqrt(1 - (x**2 + y**2)/self.R**2) )
if self.wavefrontError:
_sagPhase = self.sag
_t_m = np.exp(1j * _k * 2 * (_dz + _sagPhase) )
else:
_t_m = np.exp(1j * _k * 2 * _dz)
# Masking out the parts that fall outside the mirror clear aperture
dist_from_center = np.sqrt(x**2 + y**2)
_mirrorMask = (dist_from_center <= (self.CA/2)).astype(int)
_t_m *= _mirrorMask
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 DEFINED MIRROR'
def phase_delay(self, l, **kwargs):
_k = 2 * np.pi / l
if self.wavefrontError:
_t_m = np.exp(1j * _k * 2 * ( self.z + self.sag) )
else:
_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
def phase_delay(self, x, y, l, **kwargs):
return np.ones(np.shape(x))
class ADCObject:
def __init__(self, wedge=(1.6554, 6.28927, -1.48306)):
self.p1 = Prism('S-LAH71', 15, wedge[0], wedge[1], CA=125, name='MCD-ADC-P1')
self.p2 = Prism('S-FPL51', 15, wedge[1], wedge[2], CA=125, name='MCD-ADC-P2')
self.p3 = Prism('S-FPL51', 15, wedge[2], wedge[1], CA=125, name='MCD-ADC-P3')
self.p4 = Prism('S-LAH71', 15, wedge[1], wedge[0], CA=125, name='MCD-ADC-P4')
self.theta = None
self.set_orientation(90)
def set_orientation(self, theta):
self.p1.set_prism_orientation(theta)
self.p2.set_prism_orientation(theta)
self.p3.set_prism_orientation(180-1*theta)
self.p4.set_prism_orientation(180-1*theta)
self.theta = theta
def set_field_rotation(self, FR):
self.p1.set_prism_orientation(self.theta-FR)
self.p2.set_prism_orientation(self.theta-FR)
self.p3.set_prism_orientation(180-1*self.theta-FR)
self.p4.set_prism_orientation(180-1*self.theta-FR)
def get_phase_delay(self, x, y, l, theta=None, paraxial=False, prism=None):
if theta is not None:
self.set_orientation(theta)
else:
theta = self.theta
if not prism:
_t_p1 = self.p1.phase_delay(x, y, l)#, theta=theta, paraxial=paraxial)
_t_p2 = self.p2.phase_delay(x, y, l)#, theta=theta, paraxial=paraxial)
_t_p3 = self.p3.phase_delay(x, y, l)#, theta=180-theta, paraxial=paraxial)
_t_p4 = self.p4.phase_delay(x, y, l)#, theta=180-theta, paraxial=paraxial)
return _t_p1, _t_p2, _t_p3, _t_p4
elif prism == 1:
_t_p1 = self.p1.phase_delay(x, y, l)#, theta=theta, paraxial=paraxial)
return _t_p1
elif prism == 2:
_t_p2 = self.p2.phase_delay(x, y, l)#, theta=theta, paraxial=paraxial)
return _t_p2
elif prism == 3:
_t_p3 = self.p3.phase_delay(x, y, l)#, theta=180-1*theta, paraxial=paraxial)
return _t_p3
elif prism == 4:
_t_p4 = self.p4.phase_delay(x, y, l)#, theta=180-1*theta, paraxial=paraxial)
return _t_p4
else:
print('This prism ({}) does not exist.'.format(prism))
def add_to_optical_train(self, opticalTrainObject, d, paraxial=True):
'''
Adds the whole ADC object to an optical train object.
Because the ADC consists of multiple prisms, where each
prism has two surfaces, this is a bit more involved than
just doing opticalTrainObjec.add_object(ADC, dist=d)
Input:
opticalTrainObject - optical train object, see class OpticalTrain
d - distance in mm to next object after the ADC
paraxial - [bool] if true, each ADC only acts as a phase
transformation. No diffraction will happen
inside the prism.
'''
if not paraxial:
print('This feature has not been implemented yet. Please use paraxial mode.')
pass
if paraxial:
opticalTrainObject.add_object(self.p1, dist=d)
opticalTrainObject.add_object(self.p2, dist=1)
opticalTrainObject.add_object(self.p3, dist=8)
opticalTrainObject.add_object(self.p4, dist=1)
def calculate_dispersion(self, l1, l2, hx=0, hy=0, unit='millimeter'):
'''
Calculates the dispersion for the ADC (using geometric optics analysis)
at its current position, for a normalized on sky field angle (hx, hy).
Default field position is in the center of the FoV.
Input:
l1, l2 - Wavelength range of interest (in millimeter)
hx, hy - Field of view of the instrument (normalized between -1 and 1)
Output:
dispX, dispY - The dispersion in sky angles in milli arc seconds.
This particular ADC model is only valid for an ADC located in the pupil of the instrument.
Change the pupil ratio (Diameter_instrument_pupil / Diameter_Entrance_Pupil) accordingly
'''
if unit == 'millimeter':
l1 *= 1e3
l2 *= 1e3
elif unit == 'micron':
pass
_prismList = [self.p1, self.p2, self.p3, self.p4]
_d1x_list, _d1y_list = [], []
# _pupilRatio = 0.07943/38.542
_pupilRatio = 0.0801/38.542
hx = hx * 7.5e-3
hy = hy * 7.5e-3
ix = hx * 1/_pupilRatio
iy = hy * 1/_pupilRatio
for l in [l1,l2]:
_deltaListX = np.zeros(len(_prismList))
_deltaListY = np.zeros(len(_prismList))
for p in range(4):
_n_g = _prismList[p].find_refractive_index(l, unit='micron')
_b1x = _prismList[p].tilt1x
_b1y = _prismList[p].tilt1y
_b2x = _prismList[p].tilt2x
_b2y = _prismList[p].tilt2y
if p==0:
_t_iz_n_x = ix # Zero before first prism
_t_iz_n_y = iy # Zero before first prism
# X direction
_t_i1_n_x = _t_iz_n_x #+ _b1x
_d1x = self.calculate_deviation(_t_i1_n_x, _b1x, _b2x, _n_g)
_t_tz_n_x = _t_iz_n_x - _d1x
_t_iz_n_x = _t_tz_n_x # Set output to new input (for the next prism)
# Y direction
_t_i1_n_y = _t_iz_n_y #+ _b1y
_d1y = self.calculate_deviation(_t_i1_n_y, _b1y, _b2y, _n_g)
_t_tz_n_y = _t_iz_n_y - _d1y
_t_iz_n_y = _t_tz_n_y # Set output to new input (for the next prism)
# Adding the deviations to the list with delta terms
_deltaListX[p] = _d1x
_deltaListY[p] = _d1y
_d1x_list.append(_deltaListX.sum())
_d1y_list.append(_deltaListY.sum())
self.dispX = ( _d1x_list[1] - _d1x_list[0] ) * 3600 * 1000 * _pupilRatio
self.dispY = ( _d1y_list[1] - _d1y_list[0] ) * 3600 * 1000 * _pupilRatio
return self.dispX, self.dispY
@staticmethod
def calculate_deviation(t_i1, b1, b2, n):
'''
Returns deviation of the ray after the prism with apex angle alpha
only now using Snell's law only. Using the formula from Hagen & Tkazcyk (2011)
Units: degrees
'''
t_i1 = np.deg2rad(t_i1)
b1 = np.deg2rad(b1)
b2 = np.deg2rad(b2)
alpha = b2 - b1
delta = t_i1 - b2 - np.arcsin( n * np.sin( np.arcsin(1/n * np.sin(t_i1 - b1) ) - alpha ) )
return np.degrees(delta)
def find_optimum_angle(self, inputDispersion, l1, l2, x0=50):
'''
Simple minimizer to find the (local) minimum for a best position of the ADC.
'''
_solver = fsolve(self.optimizer, x0=x0, args=(self, inputDispersion, l1, l2))
return _solver[0]
@staticmethod
def optimizer(theta, adcObject, inputDispersion, l1, l2, FR=None):
'''
Minimization function to find the optimum rotation angle
Input:
theta - Angle in degrees relative to initial position
adcObject - the ADC class
inputDispersion - The amount of dispersion to correct for (in mas)
l1, l2 - The wavelength range to correct for (micrometer)
FR - Field rotation, not used currently.
'''
if FR:
FRdeg = np.deg2rad(FR)
inputDispersionX = np.cos(FRdeg) * inputDispersion
inputDispersionY = np.sin(FRdeg) * inputDispersion
adcObject.set_orientation(float(theta))
ADCdisp = adcObject.calculate_dispersion(l1,l2)
return np.sqrt((ADCdisp[1] - inputDispersionY)**2 + (ADCdisp[0] - inputDispersionX)**2)
else:
adcObject.set_orientation(float(theta))
ADCdisp = adcObject.calculate_dispersion(l1,l2)
ADCdisp = ADCdisp[1]
return ADCdisp - inputDispersion
class Prism:
def __init__(self, glass, thickness, tilt1, tilt2, CA, name=None):
if name:
assert type(name)==str, 'No string for a name.'
self.name = name
else:
self.name = str(self)
self.type = 'PRISM'
self.glass = glass
self.thickness = thickness
self.tilt1x = tilt1
self.tilt1y = 0
self.tilt2x = tilt2
self.tilt2y = 0
self.CA = CA
self.sag_combined = None
self.sag1 = None
self.sag2 = None
self.wavefrontError = False
self.thetaPrism = 0
# Remember original tilts for reset:
self.tilt1x_OG = tilt1
self.tilt1y_OG = 0
self.tilt2x_OG = tilt2
self.tilt2y_OG = 0
def find_refractive_index(self, l, unit='millimeter'):
if unit == 'millimeter':
l *= 1e3
elif unit == 'micron':
pass
if self.glass == 'S-FPL51':
a = [1.170105050E+00, 4.757107830E-02, 7.638324450E-01]
b = [6.162039240E-03, 2.633728760E-02, 1.418826420E+02]
elif self.glass == 'S-LAH71':
a = [1.982800310E+00, 3.167584500E-01, 2.444726460E+00]
b = [1.189874560E-02, 5.271560010E-02, 2.132206970E+02]
return self.calculate_refractive_index(l, a, b, unit='micron')
def calculate_refractive_index(self, l, a, b, unit='millimeter'):
# l in mm needs to be converted to l in micron
if unit == 'millimeter':
l *= 1e3
elif unit == 'micron':
pass
n = np.sqrt( (a[0]* l**2)/(l**2 - b[0]) + (a[1] * l**2) / (l**2 - b[1]) + (a[2] * l**2) / (l**2 - b[2]) + 1)
return n
def set_prism_orientation(self, theta):
self.tilt1x, self.tilt1y = self.tilt_rotation(self.tilt1x_OG, self.tilt1y_OG, theta)
self.tilt2x, self.tilt2y = self.tilt_rotation(self.tilt2x_OG, self.tilt2y_OG, theta)
self.thetaPrism = theta
def tilt_rotation(self, tx,ty, theta):
theta = np.deg2rad(theta-90)
tx_new = tx * np.cos(theta) - ty * np.sin(theta)
ty_new = tx * np.sin(theta) + ty * np.cos(theta)
return tx_new, ty_new
def get_surface_sag(self, x, y, theta):
# self.set_prism_orientation(theta)
self.first_surface_sag = y * np.tan(np.deg2rad(self.tilt1x)) + x * np.tan(np.deg2rad(self.tilt1y))
self.second_surface_sag = y * np.tan(np.deg2rad(self.tilt2x)) + x * np.tan(np.deg2rad(self.tilt2y))
# Masking out the parts that fall outside the prism clear aperture
dist_from_center = np.sqrt(x**2 + y**2)
_prismMask = (dist_from_center <= (self.CA/2)).astype(int)
self.first_surface_sag *= _prismMask
self.second_surface_sag *= _prismMask
return self.first_surface_sag, -1*self.second_surface_sag
def add_wavefront_error(self, z, surface=0):
assert len(z.shape)==2, 'Wavefront data does not have the proper shape. Cannot continue'
assert surface in [0,1,2], 'Surface not understood. Use 0 for transmitted wavefront error, \
1 for surface 1 or 2 for surface 2.'
if surface == 0:
self.sag_combined = z
elif surface == 1:
self.sag1 = z
elif surface == 2:
self.sag2 = z
self.wavefrontError = True
def phase_delay(self, x, y, l, theta=None, paraxial=True):
'''
This function calculates the phase delay introduced by the lens.
Input:
x, y - Coordinates of the lens
l - Wavelength in millimeters
paraxial - [bool] Apply a paraxial approximation? In the
case of a prism, this means that the prism acts
only as a phase transformation. Diffraction
inside the prism will not take place.
'''
if theta is None:
theta = self.thetaPrism
_k = 2 * np.pi / l
_n = self.find_refractive_index(l)
# Retrieve the sag of the lens surfaces for coordinates x,y
_A1, _A2 = self.get_surface_sag(x,y, theta)
# Masking out the parts that fall outside the lens clear aperture
dist_from_center = np.sqrt(x**2 + y**2)
_prismMask = (dist_from_center <= (self.CA/2)).astype(int)
if paraxial:
# Calculates the phase transformation of this prism.
# Diffraction inside the prism will not take place.
_dxy = self.thickness + _A1 + _A2
_t_p = np.exp(-1j * _k * self.thickness) * np.exp(-1j * _k * (_n - 1) * _dxy)
if self.wavefrontError:
if self.sag_combined is not None:
_t_p *= np.exp(-1j * _k * (_n - 1) * self.sag_combined)
elif self.sag1 is not None:
_t_p *= np.exp(-1j * _k * (_n - 1) * self.sag1)
elif self.sag2 is not None:
_t_p *= np.exp(-1j * _k * (_n - 1) * self.sag2)
# _t_p *= _prismMask
return _t_p
if not paraxial:
# Return two different phase transformations
# Option to add surface form error (or wavefront error)
# to each interface.
_t_p_s1 = np.exp(-1j * _k * (_n - 1) * _A1)
_t_p_s2 = np.exp(-1j * _k * (_n - 1) * _A2)
if self.wavefrontError:
if self.sag1 is not None:
_t_p_s1 *= np.exp(-1j * _k * (_n - 1) * self.sag1)
if self.sag2 is not None:
_t_p_s2 *= np.exp(-1j * _k * (_n - 1) * self.sag2)
if self.sag_combined is not None and self.sag1 is None and self.sag2 is None:
_t_p_s1 *= np.exp(-1j * _k * (_n - 1) * self.sag_combined)
# _t_p_s1 *= _prismMask
# _t_p_s2 *= _prismMask
return _t_p_s1, _t_p_s2
fourierOptics/profiles.py
+73
View File
@@ -0,0 +1,73 @@
# from .__init__ import __data_dir__
import numpy as np
import os
from astropy.io import fits
from scipy import interpolate
class Profiles:
def __init__(self):
pass
def top_hat(self, r, n, m, dpix, oversample=1, dx=0, dy=0):
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-dx)**2 + (j-dy)**2)
_Anew = (dist_from_center <= (r)).astype(int)
return _Anew
def obscured_top_hat(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 spider_top_hat(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
def gaussian(self, r, n, m, dpix, oversample=1, dx=0, dy=0):
'''
Creates a grid of n x m points, following a gaussian function
where the radius r is the FWHM of the gaussian shape.
'''
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-dx)**2 + (j-dy)**2)
return np.exp(-((j-dx)**2 + (i-dy)**2) / r**2)
def ELT_pupil(self, n, m, dpix):
# Find data directory
__pkg_dir__ = os.path.dirname(os.path.realpath(__file__))
__data_dir__ = os.path.join(__pkg_dir__, 'data')
_fits_image = __data_dir__ + '/ELT_M1_pupil_profile.fits'
_ELT_pupil = fits.open(_fits_image)[0].data
_ELT_inner_diam = 38542 # all glass diameter of M1 in mm
_dpix_fits = _ELT_inner_diam / 623 # pixels on the diameter filled with 'glass'
_resx, _resy = np.shape(_ELT_pupil)
_xf = np.linspace(-_resx/2*_dpix_fits, _resx/2*_dpix_fits, _resx)
_yf = np.linspace(-_resy/2*_dpix_fits, _resy/2*_dpix_fits, _resy)
_f_pupil = interpolate.interp2d(_xf, _yf, _ELT_pupil, kind='linear')
_x = np.linspace(-n/2*dpix, n/2*dpix, n)
_y = np.linspace(-m/2*dpix, m/2*dpix, m)
_ELT_pupil = _f_pupil(_x, _y)
return _ELT_pupil
fourierOptics/psfFitter.py
+255
View File
@@ -0,0 +1,255 @@
import numpy as np
from scipy import interpolate
from scipy.optimize import leastsq
import types
import scipy.optimize as opt
from scipy.special import j1
class PSFFitter:
'''
Wrapper around scipy.optimize.curve_fit(), so that point spread functions can be fitted
using some common profiles (e.g. Gaussian, Moffat and Airy pattern profiles).
'''
def __init__(self, x, y, psfdata):
self.x = x
self.y = y
self.psfdata = psfdata
def gaussian_PSF(self, xy, amplitude, x0, y0, sigma_x, sigma_y, theta, offset):
'''
Elliptical 2D gaussian profile
Input:
xy - Combined x,y coordinates. x and y should be a meshgrid
amplitude - amplitude of the profile, i.e. maximum of the profile
x0, y0 - shift in x and y of the central peak
sigma_x - width of the profile in x
sigma_y - width of the profile in y
theta - Rotation of the elliptical gaussian profile
offset - noise floor or sky background level
'''
_x,_y = xy
_x0 = float(x0)
_y0 = float(y0)
_a = (np.cos(theta)**2)/(2*sigma_x**2) + (np.sin(theta)**2)/(2*sigma_y**2)
_b = -(np.sin(2*theta))/(4*sigma_x**2) + (np.sin(2*theta))/(4*sigma_y**2)
_c = (np.sin(theta)**2)/(2*sigma_x**2) + (np.cos(theta)**2)/(2*sigma_y**2)
_g = offset + amplitude*np.exp( - (_a*((_x-_x0)**2) + 2*_b*(_x-_x0)*(_y-_y0) + _c*((_y-_y0)**2)))
return _g.ravel()
def fit_gaussian_PSF(self, initialGuess=None):
'''
Fit a gaussian PSF profile, a Moffat PSF profile and an Airy pattern
An OptimizeWarning is given when the best fit is reached, for some reason;
see also:
https://stackoverflow.com/questions/41725377/why-isnt-curve-fit-able-to-estimate-the-covariance-of-the-parameter-if-the-pa
Input:
initialGuess - Tuple, containing (amplitude, x0, y0, sigma_x, sigma_y, theta, offset)
'''
_popt, _pcov = opt.curve_fit(self.gaussian_PSF, (self.x,self.y), self.psfdata.flatten(), p0=initialGuess)
return self.gaussian_PSF((self.x,self.y), *_popt).reshape(self.x.shape), _popt, _pcov
def moffat_PSF(self, xy, amplitude, x0, y0, fwhm1, fwhm2, theta, offset, beta):
'''
Elliptical moffat distribution, often used to fit a PSF without diffraction rings
Equation taken from http://www.aspylib.com/doc/aspylib_fitting.html
As beta goes to infinity, the moffat profile approaches a gaussian; as mentioned
here https://www.gnu.org/software/gnuastro/manual/html_node/PSF.html
Input:
xy - Combined x,y coordinates. x and y should be a meshgrid
amplitude - amplitude of the profile, i.e. maximum of the profile
x0, y0 - shift in x and y of the central peak
fwhm1, fwhm2 - width of the profile in perpendicular directions
theta - Rotation of the elliptical moffat profile
offset - noise floor or sky background level
beta - shape parameter of the moffat profile
'''
_x,_y = xy
_x0 = float(x0)
_y0 = float(y0)
_alpha_1 = fwhm1 / (2 * np.sqrt(2**(1/beta) - 1))
_alpha_2 = fwhm2 / (2 * np.sqrt(2**(1/beta) - 1))
_a = (np.cos(theta)/_alpha_1)**2 + (np.sin(theta)/_alpha_2)**2
_b = (np.sin(theta)/_alpha_1)**2 + (np.cos(theta)/_alpha_2)**2
_c = np.sin(theta) * np.cos(theta) * (1/(_alpha_1**2) - 1/(_alpha_2**2))
_m = offset + amplitude/((1 + _a*(_x-_x0)**2 + _b*(_y-_y0)**2 + 2*_c*(_x-_x0)*(_y-_y0))**beta)
return _m.ravel()
def fit_moffat_PSF(self, initialGuess=None):
'''
Fit a Moffat PSF profile
As beta goes to infinity, the moffat profile approaches a gaussian; as mentioned
here https://www.gnu.org/software/gnuastro/manual/html_node/PSF.html
This might raise a RuntimeWarning.
Input:
initialGuess - Tuple, containing (amplitude, x0, y0, fwhm1, fwhm2, theta, offset, beta)
'''
_popt, _pcov = opt.curve_fit(self.moffat_PSF, (self.x,self.y), self.psfdata.flatten(), p0=initialGuess)
return self.moffat_PSF((self.x,self.y), *_popt).reshape(self.x.shape), _popt, _pcov
def airy_PSF(self, xy, amplitude, x0, y0, offset, radius):
'''
2D Airy pattern with offset included
Input:
xy - Combined x,y coordinates. x and y should be a meshgrid
amplitude - amplitude of the profile, i.e. maximum of the profile
x0, y0 - shift in x and y of the central peak
offset - noise floor or sky background level
'''
_x,_y = xy
_x0 = float(x0)
_y0 = float(y0)
_r = np.sqrt((_x-_x0)**2 + (_y-_y0)**2)
_Rz = 1.2196698912665045
_inp = (np.pi * _r) / (radius/_Rz)
_J1 = j1(_inp)
_fraction = np.nan_to_num(2*_J1 / _inp, nan=1.0)
_a = offset + amplitude * (_fraction)**2
return _a.ravel()
def fit_airy_PSF(self, initialGuess=None):
'''
Repackaging of scipy's curve_fit to fit an Airy pattern to the data.
No arguments are necessary, but an initial guess of the parameters
is recommended.
Input:
initialGuess - Tuple, containing (amplitude, x0, y0, offset, radius)
'''
_popt, _pcov = opt.curve_fit(self.airy_PSF, (self.x,self.y), self.psfdata.flatten(), p0=initialGuess)
return self.airy_PSF((self.x,self.y), *_popt).reshape(self.x.shape), _popt, _pcov
def reference_PSF(self, p, xy, perturbedPSF):
'''
Residual function used for the mimizer functions in
self.fit_custom_field(), see function below.
'''
_x0, _y0 = p
_x, _y = xy
_a = self.f_Eref_real(_x - _x0, _y - _y0) + 1j * self.f_Eref_imag(_x - _x0, _y - _y0)
_coupling = self.get_coupling(_a, perturbedPSF)
return 1 - np.abs(_coupling)**2
def fit_custom_field(self, Eref, Eper=None, method='local', niter_success=10, bounds=None, return_coupling=False):
'''
Fit any custom PSF to the data (self.psfdata). It requires an input
of the custom PSF in the same shape as the data. Then a global
optimization algorithm (scipy.optimization.differential_evolution)
is performed to find the translation parameters; x0 and y0.
Note:
The current implementation is not able to find other parameters,
such as amplitude, amplitude offset or radius. It assumes that
the custom PSF and data are quite similar in shape.
Also note that what is here called the PSF is actually the complex
field distribution of the PSF, i.e. not the intensity, but the
amplitude and phase information of the PSF.
Input:
Eref - Custom PSF, including phase information
Eper - optional, perturbed PSF. By default self.psfdata is taken.
method - Optimization strategy to use. Choose between
'local'
'global'
'local' first applies a local search and is very sensitive
to the initial guess.
'global' searches a global minimum. This is less sensitive
to the initial guess, but requires much more computing
time.
niter_succes- From scipy.optimize.basinhopping:
Stop the run if the global minimum candidate remains the
same for this number of iterations.
bounds - Bounds for the local minimizer
Returns:
_best_match - Best matched PSF amplitude and phase data.
_popt - Optimum parameters, i.e. translation of the centroid.
_pvar - Variances of the parameters.
_coupling - Optional; coupling factor
'''
assert Eref.shape == self.psfdata.shape, 'The shape of the new (Eper) and reference PSF (Eref) do not match.'
# Set x and y coordinates for interpolation. No 2D meshgrid is necessary.
_x = self.x[0]
_y = self.y[:,0]
# Check for custom psf data and custom bounds
if Eper is None:
Eper = self.psfdata
print('taken self.psfdata')
# Create interpolation function to match amplitude and phase to x and y coordinates.
self.f_Eref_real = interpolate.interp2d( _x, _y, Eref.real, kind='cubic')
self.f_Eref_imag = interpolate.interp2d( _x, _y, Eref.imag, kind='cubic')
# Determining a good initial guess
# Reasonable to assume >0.9 of the maximum amplitude is near the centroid position.
_max_amplitude = np.max(np.abs(Eref))
_index_i, _index_j = np.where(np.abs(Eref) > (0.9 * _max_amplitude) )
_x0_guess = np.mean(self.x[_index_i, _index_j])
_y0_guess = np.mean(self.y[_index_i, _index_j])
_p0 = [_x0_guess, _y0_guess]
if method == 'global':
# Arguments for basinhopping function (i.e. global search function)
_min_kwargs = {"args":((_x,_y), Eper), "jac":False}
_popt = opt.basinhopping(self.reference_PSF, _p0, minimizer_kwargs=_min_kwargs, niter_success=niter_success)
_p0 = _popt.x # Set optimal parameters as new first guess.
method = 'local' # Switch to local minimzer to also calculate the uncertainties
if method == 'local':
# Mimizer function using L-BFGS-B method.
if bounds is None:
bounds = [(_x.min(), _x.max()),(_y.min(), _y.max())]
_res = opt.minimize(self.reference_PSF, _p0, bounds=bounds, args=((_x,_y), Eper) )
_popt = _res.x
# Getting the uncertainty in the minimum
_ftol = 2.220446049250313e-09 # ftol assuming L-BFGS-B method in the minimizer.
_hess_inv_i = np.diag(_res.hess_inv.todense())
_pvar = max(1, abs(_res.fun)) * _ftol * _hess_inv_i
# Centroid coordinates
_x0_opt, _y0_opt = _popt
_best_match = self.f_Eref_real(_x - _x0_opt, _y - _y0_opt) + 1j * self.f_Eref_imag(_x - _x0_opt, _y - _y0_opt)
if return_coupling:
_coupling = self.get_coupling(_best_match, Eper)
return _best_match, _popt, _pvar, _coupling
else:
return _best_match, _popt, _pvar
def get_coupling(self, Eref, Eper):
'''
Calculates the coupling integral as given by eq. 6.5 in
Quasi-optical systems by Goldsmith.
This complex value is a measure of the match between a reference
field and a perturbed field. If the complex coupling factor is c,
then |c|**2 is the degradation of the perturbed field relative to
the reference field, not unlike the Strehl ratio.
Parameters:
Eref - Array; Reference field
Eper - Array; Perturbed field
Returns:
_c - Complex float; Coupling factor
'''
_c = np.sum(Eref * np.conj(Eper))/ np.sqrt(np.sum(np.abs(Eref)**2)*np.sum(np.abs(Eper)**2))
return _c
fourierOptics/wavefrontErrors.py
+506
View File
@@ -0,0 +1,506 @@
import numpy as np
import pyfftw
import types
from math import ceil
from scipy.special import factorial, j1
class WavefrontError:
'''
This class contains methods to describe wavefront errors;
either by means of zernike polynomials (description and fits)
or by means of Power Spectral Density analysis.
'''
def __init__(self, n, m):
self.n = n # Resolution in x-direction
self.m = m # Resolution in y-direction
self.resolution = (n,m)
def calculate_rnm(self,n,m,rho):
'''
@author: Tibor Agocs
Return an array with the zernike Rnm polynomial calculated at rho points.
**ARGUMENTS:**
=== ==========================================
n n order of the Zernike polynomial
m m order of the Zernike polynomial
rho Matrix containing the radial coordinates.
=== ==========================================
.. note:: For rho>1 the returned value is 0
.. note:: Values for rho<0 are silently returned as rho=0
'''
if (n-m)%2!=0:
raise Exception('n-m must be even')
if abs(m)>n:
raise Exception('The following must be true |m|<=n')
if(n==0 and m==0):
data=np.ones(rho.shape)
return data
_mask = rho > 1
rho[_mask] = 0
_Rnm = np.zeros(rho.shape)
_S = (n-np.abs(m))/2
for s in range (0,int(_S+1)):
_CR =(-1)**s*factorial(n-s)/ \
(factorial(s)*factorial(-s+(n+np.abs(m))/2)* \
factorial(-s+(n-np.abs(m))/2))
_p = _CR * ( rho**(n-2*s) )
_Rnm = _Rnm + _p
_Rnm[_mask] = 0
return _Rnm
def zernike_polynomial_polar(self,n,m,rho,theta):
'''
@author: Tibor Agocs
Returns the an array with the Zernike polynomial evaluated in the rho and
theta.
**ARGUMENTS:**
===== ==========================================
n n order of the Zernike polynomial
m m order of the Zernike polynomial
rho Matrix containing the radial coordinates.
theta Matrix containing the angular coordinates
===== ==========================================
.. note:: For rho>1 the returned value is 0
.. note:: Values for rho<0 are silently returned as rho=0
**RETURNS**
===== ===========================================
_Zmn Array with values of the Zernike polynomial
===== ===========================================
'''
_Rnm = self.calculate_rnm(n,m,rho)
# Normalization constant
_delta_m0 = bool(m==0)
_NC = np.sqrt(2*(n+1)/(1+_delta_m0))
if m >= 0:
_Zmn = _NC * _Rnm * np.cos(m*theta)
elif m < 0:
_Zmn = _NC * _Rnm * np.sin(m*theta)
return _Zmn
def zernike_polar(self, coefficient, rho, theta, noll=False, rescale='auto'):
'''
Returns the full surface created from a list of coefficients,
using the Zernike Polynomials.
rho must be between 0 and 1, if not it will be rescaled,
see also zernikePolynomialPolar()
**ARGUMENTS:**
============ ==========================================
coefficient n order of the Zernike polynomial
rho Matrix containing the radial coordinates.
theta Matrix containing the angular coordinates
noll Boolean to indicate to use noll indices
============ ==========================================
**RETURNS**
**RETURNS**
===== ===========================================
_zP Array with values of the Zernike polynomial
===== ===========================================
'''
if rescale == 'auto':
if np.max(rho) > np.sqrt(2):
print('Rescaling!')
rho = np.sqrt(2) * rho / np.max(rho)
_termList = []
for i in range(len(coefficient)):
_C = coefficient[i]
if noll:
_jTerm = self.convert_noll_to_standard(i+1)
_n,_m = self.i_to_nm(i)
_termList.append( _C * self.zernike_polynomial_polar(_n,_m,rho,theta) )
_zP = sum(_termList)
return _zP
def zernike_cartesian(self, coefficient, x, y, noll=False, addToData=False, **kwargs):
'''
Converts the cartesian grid x,y into polar
coordinates and returns the data, such that
it can be plotted with e.g. plt.imshow()
**ARGUMENTS:**
============ ==========================================
coefficient n order of the Zernike polynomial
rho Matrix containing the radial coordinates.
theta Matrix containing the angular coordinates
noll Boolean to indicate to use noll indices
============ ==========================================
**RETURNS**
===== ===========================================
_zP Array with values of the Zernike polynomial
===== ===========================================
'''
_x0 = (np.max(x) - np.min(x)) / 2 + np.min(x)
_y0 = (np.max(y) - np.min(y)) / 2 + np.min(y)
_rho = np.sqrt(( x - _x0 )**2 + ( y - _y0 )**2 )
_theta = np.arctan2( ( y - _y0 ) , ( x - _x0 ) )
_zP = self.zernike_polar(coefficient, _rho, _theta, noll, **kwargs)
if not addToData:
self.data = _zP
else:
self.data += _zP
self.rms = np.sqrt(np.mean(np.asarray(self.data)**2))
self.ptv = np.abs(np.max(self.data)) + np.abs(np.min(self.data))
return _zP
def i_to_nm(self, i):
'''
Converts a single index to n and m for the Zernike Polynomials
'''
n = ceil((-3+np.sqrt(9+8*i))/2)
m = 2 * i - n*(n+2)
return int(n),int(m)
def convert_noll_to_standard(self, j):
'''
'''
conversionList = [1, 3, 2, 5, 4, 6, 9, 7, 8, 10, 15, 13, 11, 12, 14, 21, 19, 17, 16, 18, 20, 27, 25, 23, 22, 24, 26, 28, 35, 33, 31, 29, 30, 32, 34, 36, 45, 43, 41, 39, 37, 38, 40, 42, 44, 55, 53, 51, 49, 47, 46, 48, 50, 52, 54, 65, 63, 61, 59, 57, 56, 58, 60, 62, 64, 66, 77, 75, 73, 71, 69, 67, 68, 70, 72, 74, 76, 78, 91, 89, 87, 85, 83, 81, 79, 80, 82, 84, 86, 88, 90, 105, 103, 101, 99, 97, 95, 93, 92, 94, 96, 98, 100, 102, 104, 119, 117, 115, 113, 111, 109, 107, 106, 108, 110, 112, 114, 116, 118, 120]
i = conversionList.index(j)
return i
def fit_zernike_coefficients(self, data, x=None, y=None, modes=21):
'''
Fits the zernike coefficients to a data set, using
that the zernike modes can be treated independently.
The modes are set following standard index notation.
We solve by an analytical least squares fit
Z * A = F
Z.T * Z * A = Z.T * F
A = (Z.T * Z)^-1 * Z.T * F
**ARGUMENTS:**
===== ==========================================
data array; the data that needs fitting
x array; the x-coordinates of the data
y array; the y-coordinates of the data
modes integer; number of modes to be fitted
===== ==========================================
**RETURNS**
===== ===========================================
_A Array with values of the Zernike polynomial
===== ===========================================
'''
if x is None and y is None:
y,x = np.mgrid[-1:1:self.n*1j, -1:1:self.m*1j]
_F = data.flatten()[:,np.newaxis]
_x = x.flatten()
_y = y.flatten()
_x0 = (np.max(_x) - np.min(_x)) / 2 + np.min(_x)
_y0 = (np.max(_y) - np.min(_y)) / 2 + np.min(_y)
_rho = np.sqrt( ( _x - _x0 )**2 + ( _y - _y0 )**2 )
_theta = np.arctan2( ( _y - _y0 ), ( _x - _x0 ) )
_Zern = []
for i in range(modes):
_n, _m = self.i_to_nm(i)
_Znm = self.zernike_polynomial_polar(_n,_m,_rho,_theta)
_Zern.append(_Znm)
_Znm = None
_Zern = np.asarray(_Zern).T
_Zern_transpose = _Zern.T
_pseudo_inv = np.linalg.pinv(np.dot(_Zern_transpose, _Zern))
_int1 = np.dot(_pseudo_inv, _Zern_transpose)
_A = np.dot(_int1,_F)
self.zernikeCoefficients = _A[:,0]
return _A[:,0]
def scale_zernike_coefficients(self, c, rms):
# Scaling the coefficients to some total value
_rms_coeffic = np.sqrt(np.sum( np.asarray(c)**2))
_scale_factor = rms / _rms_coeffic
_c = (_scale_factor * np.asarray(c)).tolist()
# _rms_coeffic = np.sqrt(np.sum( np.asarray(_c)**2))
self.zernikeCoefficients = _c
return _c
def power_spectral_density(self, s0, p, rho_c, dx, dy, A):
# Lorentzian equation used as PSD function
m = self.m
n = self.n
UmList = [(M - ( m + 1) / 2) / (m * dx) for M in range(m)]
VnList = [(N - ( n + 1) / 2) / (n * dy) for N in range(n)]
VnList, UmList = np.meshgrid(UmList, VnList)
_rho_mn_list = np.sqrt(UmList**2 + VnList**2)
_h0list = 1 / ( 1 + np.power( ( _rho_mn_list / rho_c ), p) )
_h0_norm = np.sum(_h0list)
_rho_mn_list = np.asarray(_rho_mn_list).reshape(m,n)
_PSD = (s0**2 * A) / (_h0_norm * (1 + np.power((_rho_mn_list/rho_c), p)))
self.PSD2D = _PSD
return _PSD, _rho_mn_list
def phase_map(self, phase):
# Uniform distribution of the phase
_phase_in = pyfftw.empty_aligned((self.n,self.m), dtype='complex128')
_phase_in[:] = phase
_phase_out = pyfftw.empty_aligned((self.n, self.m), dtype='complex128')
_FFTW_Obj = pyfftw.FFTW( _phase_in, _phase_out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
_FFTW_Obj()
_phaseAngle = np.angle(_phase_out)
return _phaseAngle
def generate_height_map(self, s0, p, rho_c, addToData=False):
'''
This function generates a surface height distribution based
on a Lorentzian PSD profile. Following the paper by E. Sidick (2009).
2020-01-20: Removed "mean=1e-4" from function def. --> Did this do anything?
'''
# Calculation of the necessary constants
_dx = 1 / self.m
_dy = 1 / self.n
_A = (self.m * _dx) * (self.n * _dy)
_phase = np.random.uniform(0,1,(self.m,self.m))
# Calculating the PSD and phase
_PSD, _rho_mn_list = self.power_spectral_density(s0, p, rho_c, _dx, _dy, _A)
_phiMap = self.phase_map(_phase)
# Shifting the coordinates for the FFT
_PSD = np.fft.ifftshift(_PSD)
# Doing the fourier transform
_in = pyfftw.empty_aligned((self.n,self.m), dtype='complex128')
_in[:] = np.exp( 1j * _phiMap) * np.sqrt(_A * np.abs(_PSD))
_out = pyfftw.empty_aligned((self.n,self.m), dtype='complex128')
_FFTW_Obj = pyfftw.FFTW( _in, _out, axes=(0,1), direction='FFTW_BACKWARD', threads=6, flags=('FFTW_ESTIMATE',))
_FFTW_Obj()
# The final height distribution
_hMap = 1 / ( _dx * _dy ) * _out
if not addToData:
self.data = _hMap.real
else:
self.data += _hMap.real
self.rms = np.sqrt(np.mean(np.asarray(self.data)**2))
self.ptv = np.abs(np.max(self.data)) + np.abs(np.min(self.data))
return self.data
def scale_rms(self, factor, data=None, filteredDataFlag=False):
# Scale the surface data such that the RMS gets multiplied by factor
_factor = factor ** 2
if data is None:
if filteredDataFlag:
self.scaledPSD2D = self.filteredPSD2D * _factor
_data = self.filteredData
else:
self.calculate_2D_PSD()
self.scaledPSD2D = self.PSD2D * _factor
_data = self.data
elif data is not None:
self.PSD2D = self.calculate_2D_PSD(data=data)
self.scaledPSD2D = self.PSD2D * _factor
_data = data
# Fourier transform of the surface profile
_in = pyfftw.empty_aligned((self.n,self.m), dtype='complex128')
_in[:] = _data
_out = pyfftw.empty_aligned((self.n,self.m), dtype='complex128')
_FFTW_Obj = pyfftw.FFTW( _in, _out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
_FFTW_Obj()
_phase = np.angle(_out)
_in = pyfftw.empty_aligned((self.n,self.m), dtype='complex128')
_in[:] = np.exp(1j*_phase) * np.sqrt(np.fft.fftshift(self.scaledPSD2D))
_out = pyfftw.empty_aligned((self.n,self.m), dtype='complex128')
_FFTW_Obj = pyfftw.FFTW( _in, _out, axes=(0,1), direction='FFTW_BACKWARD', threads=6, flags=('FFTW_ESTIMATE',))
_FFTW_Obj()
self.scaledData = _out.real
self.scaledrms = np.sqrt(np.mean(np.asarray(self.scaledData)**2))
self.scaledptv = np.abs(np.max(self.scaledData)) + np.abs(np.min(self.scaledData))
return _out.real
def calculate_2D_PSD(self, data=None):
''' Returns the 2D PSD map '''
if data is None:
_in = pyfftw.empty_aligned((self.n,self.m), dtype='complex128')
_in[:] = self.data
_out = pyfftw.empty_aligned((self.n,self.m), dtype='complex128')
elif data is not None:
_n, _m = np.asarray(data).shape
_in = pyfftw.empty_aligned((_n,_m), dtype='complex128')
_in[:] = np.asarray(data)
_out = pyfftw.empty_aligned((_n,_m), dtype='complex128')
_FFTW_Obj = pyfftw.FFTW( _in, _out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
_FFTW_Obj()
if data is None:
self.D1 = np.fft.fftshift(_out)
self.PSD2D = np.abs(self.D1)**2
return self.PSD2D
elif data is not None:
_D1 = np.fft.fftshift(_out)
_PSD2D = np.abs(_D1)**2
return _PSD2D
def calculate_1D_PSD(self, data=None, filteredDataFlag=False, rmax=None):
''' Returns the azimuthally averaged PSD map '''
if data is None:
try:
if not filteredDataFlag:
self.PSD1D = self.azimuthal_average(rmax)
else:
self.filteredPSD1D = self.azimuthal_average(filteredDataFlag=filteredDataFlag)
return self.filteredPSD1D
except AttributeError:
self.calculate_2D_PSD(data)
self.PSD1D = self.azimuthal_average(rmax)
return self.PSD1D
elif data is not None:
_PSD2D = self.calculate_2D_PSD(data=data)
_PSD1D = self.azimuthal_average(PSD=_PSD2D, rmax=rmax)
return _PSD1D
def set_lowpass_filter(self, r, dpix=1):
''' Returns a map where all frequencies higher than r have been filtered out '''
self.calculate_2D_PSD()
# Masking out the parts that fall outside the lens clear aperture
_x, _y = np.mgrid[-self.n/2*dpix : self.n/2*dpix : self.n*1j, -self.m/2*dpix : self.m/2*dpix : self.m*1j]
_dist_from_center = np.sqrt(_x**2 + _y**2) * dpix
_mask = (_dist_from_center <= (r)).astype(int)
_maskedPSD2D = self.PSD2D.copy()
_maskedPSD2D *= _mask
_phase = np.angle(np.fft.fft2(self.data))
self.filteredPSD2D = _maskedPSD2D
self.filteredData = np.fft.ifft2(np.exp(1j*_phase) * np.sqrt(np.fft.fftshift(self.filteredPSD2D))).real
return self.filteredData
def set_highpass_filter(self, r, dpix=1):
'''Returns a map where all spatial frequencies lower than r have been filtered out'''
self.calculate_2D_PSD()
# Masking out the parts that fall outside the lens clear aperture
# x,y = np.ogrid[-self.resolution[0]/2:self.resolution[0]/2, -self.resolution[1]/2:self.resolution[1]/2]
_x, _y = np.mgrid[-self.n/2*dpix : self.n/2*dpix : self.n*1j, -self.m/2*dpix : self.m/2*dpix : self.m*1j]
_dist_from_center = np.sqrt(_x**2 + _y**2) * dpix
_mask = (_dist_from_center <= (r)).astype(int)
_maskedPSD2D = self.PSD2D.copy()
_maskedPSD2D *= (1-_mask)
_phase = np.angle(np.fft.fft2(self.data))
self.filteredPSD2D = _maskedPSD2D
self.filteredData = np.fft.ifft2(np.exp(1j*_phase) * np.sqrt(np.fft.fftshift(self.filteredPSD2D))).real
return self.filteredData
def set_bandpass_filter(self, rmin, rmax, dpix=1):
'''Returns a map where all spatial frequencies lower than r2 and higher than r1 have been filtered out'''
assert rmax > rmin, 'rmax < rmin, not physical!'
self.calculate_2D_PSD()
# Masking out the parts that fall outside the lens clear aperture
# x,y = np.ogrid[-self.resolution[0]/2:self.resolution[0]/2, -self.resolution[1]/2:self.resolution[1]/2]
_x, _y = np.mgrid[-self.n/2*dpix : self.n/2*dpix : self.n*1j, -self.m/2*dpix : self.m/2*dpix : self.m*1j]
_dist_from_center = np.sqrt(_x**2 + _y**2) * dpix
_mask_1 = (_dist_from_center <= (rmax)).astype(int)
_mask_2 = (_dist_from_center >= (rmin)).astype(int)
_maskedPSD2D = self.PSD2D.copy()
_maskedPSD2D *= (_mask_1 & _mask_2)
_phase = np.angle(np.fft.fft2(self.data))
self.filteredPSD2D = _maskedPSD2D
self.filteredData = np.fft.ifft2(np.exp(1j*_phase) * np.sqrt(np.fft.fftshift(self.filteredPSD2D))).real
return self.filteredData
def azimuthal_average(self, PSD=None, center=None, filteredDataFlag=False, rmax=None):
"""
Calculate the azimuthally averaged radial profile.
image - The 2D image
center - The [x,y] pixel coordinates used as the center. The default is
None, which then uses the center of the image (including
fracitonal pixels).
"""
# Calculate the indices from the image
_y, _x = np.indices(self.resolution)
if not center:
center = np.array([(_x.max()-_x.min())/2.0, (_y.max()-_y.min())/2.0])
_r = np.hypot(_x - center[0], _y - center[1])
# Get sorted radii
_ind = np.argsort(_r.flat)
_r_sorted = _r.flat[_ind]
if not filteredDataFlag and PSD is None:
_PSD = self.PSD2D
elif PSD is None:
_PSD = self.filteredPSD2D
elif PSD is not None:
_PSD = PSD
if rmax == True:
_PSD[ _r > _x.max()/2 ] = np.nan
elif rmax:
_PSD[ _r > rmax ] = np.nan
_i_sorted = _PSD.flat[_ind]
# Get the integer part of the radii (bin size = 1)
_r_int = _r_sorted.astype(int)
# Find all pixels that fall within each radial bin.
_deltar = _r_int[1:] - _r_int[:-1] # Assumes all radii represented
_rind = np.where(_deltar)[0] # location of changed radius
_nr = _rind[1:] - _rind[:-1] # number of radius bin
# Cumulative sum to figure out sums for each radius bin
_csim = np.cumsum(_i_sorted, dtype=float)
_tbin = _csim[_rind[1:]] - _csim[_rind[:-1]]
_radial_prof = _tbin / _nr
return _radial_prof
def generate_sine_pattern(self, ampx, ampy, perx, pery):
_x, _y = np.mgrid[-1:1:self.n*1j, -1:1:self.m*1j]
_z = ampx * np.sin(2*np.pi*_x*perx/2) + ampy * np.sin(2 * np.pi * _y * pery/2)
return _z
def generate_tilt(self, x, y, tx, ty):
_z = y * np.tan(np.deg2rad(tx)) + x * np.tan(np.deg2rad(ty))
return _z
test cases (standalone)/fourierOptics_01_diffracting_aperture.py
+735
View File
@@ -0,0 +1,735 @@
#!/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', transparent=True)
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, _y, _E_out.real, kind='linear')
_f_ip_imag = interpolate.interp2d(_x, _y, _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
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 phaseDelayFunction(self, x, y, l, paraxial=False):
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
if paraxial:
_dxy = self.thickness - (x**2 + y**2)/2 * (1/self.R1 - 1/self.R2)
_k = 2 * np.pi / 30e-3
_n = self.findRefractiveIndex(self.referenceWavelength)
_t_l = np.exp(-1j * _k * self.thickness)* np.exp(-1j * _k * (_n - 1) * _dxy)
# 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 phaseDelayFunction(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 phaseDelayFunction(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 phaseDelayFunction(self, l):
_k = 2 * np.pi / l
_t_m = np.exp(-1j * _k * 2 * self.z)
return _t_m
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 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)
_wf = U_in
for _lens in self.opticalTrain:
_i = self.opticalTrain.index(_lens)
_dz = self.sequentialDistance[_i]
_wf = _lens.phaseDelayFunction(x, y, l, paraxial=paraxial)
# Perform Fresnel number test:
_F = self.fresnel_number()
if _F < 5:
# Fresnel propagation
_wf, x, y = self.fresnel_propagation()
elif _F >= 5:
# Angular spectrum (p2p) propagation
_wf, x, y = self.angular_spectrum_propagation()
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/50*l
n, m = 2001, 2001
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 = profiles()
# A = prof.spiderTopHat(250, 100, 2001, 2001, dpix=1, oversample=1)
A = prof.topHat(4*l, n, m, dpix=dpix, oversample=1)
# Alternatively, create a gaussian profile
# gauss = GaussianBeam(0.3, l*1000)
# A = gauss.fundamentalMode(0,n,m,np.max(x))
# 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))
# The input distribution gets a curved wavefront after the lens
U_in = A
Fresnel = False
P2P = not Fresnel
oC = [x.min(), x.max(), y.min(), y.max()]
# Performing the propagation using Fresnel approximation to the focal plane.
if Fresnel:
savename = 'fourierOptics_01_diffracting_aperture_fresnel.pdf'
Exyz,_,_ = fou.FresnelPropagationFFTW(U_in, x, y, dz=4*l, oversample=None, show=True, save=savename, outputCoordinates=oC)
if P2P:
savename = 'fourierOptics_01_diffracting_aperture_p2p.pdf'
Exyz,_,_ = fou.planeWavePropagationFFTW(A, dz=4*l, dpix=dpix, oversample=None, show=True, save=savename)
# Plot a cross section of the new field (useful to better see the diffraction rings)
fig = plt.figure()
ax = fig.add_subplot(111)
ax.grid(True, alpha=0.3)
ax.semilogy(np.linspace(np.min(x), np.max(x),n), np.abs(Exyz[1000,:]**2), 'r-')
ax.set_xlabel('x (mm)')
ax.set_ylabel('Intensity')
ax.tick_params(axis ='both', direction='in', which='both')
if Fresnel:
plt.savefig('fourierOptics_01_diffracting_aperture_fresnel_slice.pdf', bbox_inches='tight', transparent=True)
if P2P:
plt.savefig('fourierOptics_01_diffracting_aperture_p2p_slice.pdf', bbox_inches='tight', transparent=True)
plt.show()
# The following produces a cross-section of the propagation for various z, so that you can see how the beam evolves.
if True:
Ics = []
Ics_x = []
Isum = []
xmax = []
zList = np.linspace(1,10.1,18).tolist()
xin, yin = x, y
for i in zList:
title = f"{zList.index(i)+1:03d}"
print (' Calculating {0}/{1}'.format(title, f"{len(zList):03d}"), end="\r")
# Propagate the wavefront
if Fresnel:
Exyz, x, y = fou.planeWavePropagationFFTW(A, dz=i*l, dpix=dpix, oversample=None, show=False, save=False)#, scaleInputArray=(n,m))
if P2P:
Exyz, x, y = fou.FresnelPropagationFFTW(U_in, xin, yin, dz=i*l, oversample=None, show=False, save=False, outputCoordinates=oC)
# Process the results by taking a slice of the data
Isum.append(np.sum(np.abs(Exyz)**2))
Ics.append((np.abs(Exyz[int(Exyz.shape[0]/2)])**2))
Ics_x.append(x.real[0])
y_ext = [y.real[0,0],y.real[-1,0]]
# Reset the data
x = None
y = None
# Plotting all slices
fig = plt.figure(figsize=(8,4))
ax = fig.add_subplot(211)
ax.tick_params(axis ='both', direction='in')
# Plotting the square root of the amplitude for better contrast
im = ax.imshow(np.asarray(Ics).T, cmap='hot', aspect='auto', interpolation='bicubic', extent=[zList[0],zList[-1],y_ext[0],y_ext[1]])
ax.yaxis.set_major_formatter(FormatStrFormatter('%.1E'))
ax.set_ylim([-0.2, 0.2])
# Plotting the intensity relative to the input
ax2 = fig.add_subplot(212)
ax2.plot(zList, np.asarray(Ics)[:,int(np.asarray(Ics).shape[1]/2)], 'k-')
ax2.yaxis.set_major_formatter(FormatStrFormatter('%.1E'))
ax2.set_xlabel('Propagation distance ($\\times \lambda$)')
ax2.tick_params(axis ='both', direction='in')
ax2.set_ylabel('Intensity')
ax2.set_xlim([np.min(zList), np.max(zList)])
ax2.grid(True, alpha=0.3)
# Saving the figure? Then show.
plt.tight_layout()
# plt.savefig('fresnel_cross.pdf', bbox_inches='tight')
# plt.savefig('fresnel_cross.png', bbox_inches='tight')
if P2P:
plt.savefig('fourierOptics_01_diffracting_aperture_p2p_cross.pdf', bbox_inches='tight', transparent=True)
if Fresnel:
plt.savefig('fourierOptics_01_diffracting_aperture_fresnel_cross.pdf', bbox_inches='tight', transparent=True)
plt.show()
# Closing down the script and stopping the timer.
end = time.time()
print('The program took {:.3f} seconds!'.format((end-start)))
test cases (standalone)/fourierOptics_025_single_lens_system.py
+729
View File
@@ -0,0 +1,729 @@
#!/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, _y, _E_out.real, kind='linear')
_f_ip_imag = interpolate.interp2d(_x, _y, _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
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 phaseDelayFunction(self, x, y, l, paraxial=False):
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
if paraxial:
_dxy = self.thickness - (x**2 + y**2)/2 * (1/self.R1 - 1/self.R2)
_k = 2 * np.pi / 30e-3
_n = self.findRefractiveIndex(self.referenceWavelength)
_t_l = np.exp(-1j * _k * self.thickness)* np.exp(-1j * _k * (_n - 1) * _dxy)
# 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 phaseDelayFunction(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 phaseDelayFunction(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 phaseDelayFunction(self, l):
_k = 2 * np.pi / l
_t_m = np.exp(-1j * _k * 2 * self.z)
return _t_m
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 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)
_wf = U_in
for _lens in self.opticalTrain:
_i = self.opticalTrain.index(_lens)
_dz = self.sequentialDistance[_i]
_wf = _lens.phaseDelayFunction(x, y, l, paraxial=paraxial)
# Perform Fresnel number test:
_F = self.fresnel_number()
if _F < 5:
# Fresnel propagation
_wf, x, y = self.fresnel_propagation()
elif _F >= 5:
# Angular spectrum (p2p) propagation
_wf, x, y = self.angular_spectrum_propagation()
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/50
n, m = 2001, 2001
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 = profiles()
# A = prof.spiderTopHat(250, 100, 2001, 2001, dpix=1, oversample=1)
A = prof.topHat(2.5, n, m, dpix=dpix, oversample=1)
# Alternatively, create a gaussian profile
# gauss = GaussianBeam(0.3, l*1000)
# A = gauss.fundamentalMode(0,n,m,np.max(x))
# 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 a focal length ~100 mm, with a clear aperture radius of 4 mm and thickness of 1 mm.
L1 = lensObject('N-BK7', 50, 1e99, 1, 4, name='L1')
print('Focal length of lens {} is {:.1f} mm.'.format(L1.name, L1.f))
# Calculate the phase delay introduced by the lens
t_l = L1.phaseDelayFunction(x,y,1.35e-6, paraxial=False)
# The input distribution gets a curved wavefront after the lens
U_in = A * t_l
# Performing the propagation using Fresnel approximation to the focal plane.
oC = [-5, 5, -5, 5]
savename = 'simple_lens.pdf'
Exyz,_,_ = fou.FresnelPropagationFFTW(U_in, x, y, dz=L1.f, oversample=None, show=True, save=savename, outputCoordinates=oC)
# Plot a cross section of the new field (useful to better see the diffraction rings)
fig = plt.figure()
ax = fig.add_subplot(111)
ax.grid(True, alpha=0.3)
ax.semilogy(np.linspace(np.min(oC), np.max(oC),n), np.abs(Exyz[1000,:]**2), 'r-')
ax.set_xlabel('x (mm)')
ax.set_ylabel('Intensity')
ax.tick_params(axis ='both', direction='in', which='both')
plt.savefig('simple_lens_slice.pdf', bbox_inches='tight', transparent=True)
plt.show()
# The following produces a cross-section of the propagation for various z, so that you can see how the beam evolves.
if True:
Ics = []
Ics_x = []
Isum = []
xmax = []
zList = np.linspace(20,220,20).tolist()
xin, yin = x, y
for i in zList:
title = f"{zList.index(i)+1:03d}"
print (' Calculating {0}/{1}'.format(title, f"{len(zList):03d}"), end="\r")
# Propagate the wavefront
# Exyz, x, y = fou.planeWavePropagationFFTW(A, dz=i*l, dpix=dpix, oversample=1, show=False, save=False)#, scaleInputArray=(n,m))
Exyz, x, y = fou.FresnelPropagationFFTW(U_in, xin, yin, dz=i, oversample=None, show=False, save=False, outputCoordinates=oC)
# Process the results by taking a slice of the data
Isum.append(np.sum(np.abs(Exyz)**2))
Ics.append((np.abs(Exyz[int(Exyz.shape[0]/2)])**2))
Ics_x.append(x.real[0])
y_ext = [y.real[0,0],y.real[-1,0]]
# Reset the data
x = None
y = None
# Plotting all slices
fig = plt.figure(figsize=(8,4))
ax = fig.add_subplot(211)
ax.tick_params(axis ='both', direction='in')
# Plotting the square root of the amplitude for better contrast
im = ax.imshow(np.sqrt(np.asarray(Ics)).T, cmap='hot', aspect='auto', interpolation='bicubic', extent=[zList[0],zList[-1],y_ext[0],y_ext[1]])
ax.yaxis.set_major_formatter(FormatStrFormatter('%.1E'))
# Plotting the intensity relative to the input
ax2 = fig.add_subplot(212)
ax2.plot(zList, np.asarray(Ics)[:,int(np.asarray(Ics).shape[1]/2)], 'k-')
ax2.yaxis.set_major_formatter(FormatStrFormatter('%.1E'))
ax2.set_xlabel('Propagation distance (mm)')
ax2.tick_params(axis ='both', direction='in')
ax2.set_ylabel('Amplitude')
ax2.set_xlim([np.min(zList), np.max(zList)])
ax2.grid(True, alpha=0.3)
# Saving the figure? Then show.
plt.tight_layout()
# plt.savefig('fresnel_cross.pdf', bbox_inches='tight')
# plt.savefig('fresnel_cross.png', bbox_inches='tight')
plt.savefig('simple_lens_cross.pdf', bbox_inches='tight', transparent=True)
plt.show()
# Closing down the script and stopping the timer.
end = time.time()
print('The program took {:.3f} seconds!'.format((end-start)))
test cases (standalone)/fourierOptics_02_telescope.py
+771
View File
@@ -0,0 +1,771 @@
#!/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, _y, _E_out.real, kind='linear')
_f_ip_imag = interpolate.interp2d(_x, _y, _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.sqrt(np.abs(_E_out_ip)), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()])
_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')
# _fig.colorbar(_im, ax=_ax, label='$\\sqrt{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
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 phaseDelayFunction(self, x, y, l, paraxial=False):
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
if paraxial:
_dxy = self.thickness - (x**2 + y**2)/2 * (1/self.R1 - 1/self.R2)
_k = 2 * np.pi / 30e-3
_n = self.findRefractiveIndex(self.referenceWavelength)
_t_l = np.exp(-1j * _k * self.thickness)* np.exp(-1j * _k * (_n - 1) * _dxy)
# 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 phaseDelayFunction(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 phaseDelayFunction(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
_k = 2 * np.pi / l
_dz = (x**2 + y**2) / (4 * self.f)
# Masking out the parts that fall outside the mirror clear aperture
dist_from_center = np.sqrt(x**2 + y**2)
_mirrorMask = (dist_from_center <= (self.CA/2)).astype(int)
_t_m *= _mirrorMask
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 phaseDelayFunction(self, l):
_k = 2 * np.pi / l
_t_m = np.exp(1j * _k * 2 * self.z)
return _t_m
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 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)
_wf = U_in
for _lens in self.opticalTrain:
_i = self.opticalTrain.index(_lens)
_dz = self.sequentialDistance[_i]
_wf = _lens.phaseDelayFunction(x, y, l, paraxial=paraxial)
# Perform Fresnel number test:
_F = self.fresnel_number()
if _F < 5:
# Fresnel propagation
_wf, x, y = self.fresnel_propagation()
elif _F >= 5:
# Angular spectrum (p2p) propagation
_wf, x, y = self.angular_spectrum_propagation()
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 = 8
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)
prof = profiles()
A = prof.topHat(4000, n, m, dpix=dpix, oversample=1)
# A = prof.spiderTopHat(4000, 1000, n, m, dpix=dpix, oversample=1)
# A = prof.obscuredTopHat(4000, 1000, n, m, dpix=dpix, oversample=1)
# Alternatively, create a gaussian profile
# gauss = GaussianBeam(0.3, l*1000)
# A = gauss.fundamentalMode(0,n,m,np.max(x))
# 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 mirror with a focal length 64 m, with a clear aperture radius of 4 m.
M1 = parabolicMirror(dx=0, dy=0, name='M1', CA=8000, f=64000)
print('Focal length of mirror {} is {:.1f} mm.'.format(M1.name, M1.f))
# Calculate the phase delay introduced by the mirror
t_l = M1.phaseDelayFunction(x,y,l)
# The input distribution gets a curved wavefront after the mirror
U_in = A * t_l
plt.imshow(np.abs(U_in))
plt.colorbar()
plt.savefig('telescope_mask.pdf', bbox_inches='tight', transparent=True)
plt.show(block=False)
# Performing the propagation using Fresnel approximation to the focal plane.
# oC = [-0.001, 0.001, -0.001, 0.001]
oC = [-2, 2, -2, 2]
# savename = 'fourierOptics_02_telescope_fresnel.pdf'
Exyz,xout,yout = fou.FresnelPropagationFFTW(U_in, x, y, dz=M1.f, oversample=None, show=True, save=False, outputCoordinates=oC)
# Plot a cross section of the new field (useful to better see the diffraction rings)
fig = plt.figure()
ax = fig.add_subplot(111)
ax.grid(True, alpha=0.3)
ax.semilogy(np.linspace(np.min(oC), np.max(oC),n), np.abs(Exyz[int((n-1)/2),:]**2), 'r-')
ax.set_xlabel('x (mm)')
ax.set_ylabel('Intensity')
ax.tick_params(axis ='both', direction='in', which='both')
# plt.savefig('fourierOptics_02_telescope_slice.pdf', bbox_inches='tight', transparent=True)
plt.show()
# Combined plot
# Plotting the results
_fig = plt.figure(figsize=(12,3))
_ax = _fig.add_subplot(131)
_im = _ax.imshow(np.abs(Exyz), cmap='hot', extent=oC)
_ax.set_xlim([oC[0], oC[1]])
_ax.set_ylim([oC[2], oC[3]])
_ax.set_xlabel('x [mm]')
_ax.set_ylabel('y [mm]')
_fig.colorbar(_im, ax=_ax, label='Amplitude')
_ax2 = _fig.add_subplot(132)
_im2 = _ax2.imshow(np.angle(Exyz), cmap='hot', extent=oC)
_ax2.set_xlim([oC[0], oC[1]])
_ax2.set_ylim([oC[2], oC[3]])
_ax2.set_xlabel('x [mm]')
_ax2.set_ylabel('y [mm]')
_fig.colorbar(_im2, ax=_ax2, label='Phase')
_ax3 = _fig.add_subplot(133)
_ax3.grid(True, alpha=0.3)
_ax3.semilogy(np.linspace(np.min(oC), np.max(oC),n), np.abs(Exyz[int((n-1)/2),:]**2), 'r-')
_ax3.set_xlabel('x (mm)')
_ax3.set_ylabel('Intensity')
_ax3.tick_params(axis ='both', direction='in', which='both')
plt.tight_layout()
plt.savefig('fourierOptics_02_telescope_complete.pdf', bbox_inches='tight', transparent=True)
plt.show()
# The following produces a cross-section of the propagation for various z, so that you can see how the beam evolves.
if False:
Ics = []
Ics_x = []
Isum = []
xmax = []
zList = np.linspace(20,220,20).tolist()
xin, yin = x, y
for i in zList:
title = f"{zList.index(i)+1:03d}"
print (' Calculating {0}/{1}'.format(title, f"{len(zList):03d}"), end="\r")
# Propagate the wavefront
# Exyz, x, y = fou.planeWavePropagationFFTW(A, dz=i*l, dpix=dpix, oversample=1, show=False, save=False)#, scaleInputArray=(n,m))
Exyz, x, y = fou.FresnelPropagationFFTW(U_in, xin, yin, dz=i, oversample=None, show=False, save=False, outputCoordinates=oC)
# Process the results by taking a slice of the data
Isum.append(np.sum(np.abs(Exyz)**2))
Ics.append((np.abs(Exyz[int(Exyz.shape[0]/2)])**2))
Ics_x.append(x.real[0])
y_ext = [y.real[0,0],y.real[-1,0]]
# Reset the data
x = None
y = None
# Plotting all slices
fig = plt.figure(figsize=(8,4))
ax = fig.add_subplot(211)
ax.tick_params(axis ='both', direction='in')
# Plotting the square root of the amplitude for better contrast
im = ax.imshow(np.sqrt(np.asarray(Ics)).T, cmap='hot', aspect='auto', interpolation='bicubic', extent=[zList[0],zList[-1],y_ext[0],y_ext[1]])
ax.yaxis.set_major_formatter(FormatStrFormatter('%.1E'))
# Plotting the intensity relative to the input
ax2 = fig.add_subplot(212)
ax2.plot(zList, np.asarray(Ics)[:,int(np.asarray(Ics).shape[1]/2)], 'k-')
ax2.yaxis.set_major_formatter(FormatStrFormatter('%.1E'))
ax2.set_xlabel('Propagation distance (mm)')
ax2.tick_params(axis ='both', direction='in')
ax2.set_ylabel('Intensity')
ax2.set_xlim([np.min(zList), np.max(zList)])
ax2.grid(True, alpha=0.3)
# Saving the figure? Then show.
plt.tight_layout()
# plt.savefig('fresnel_cross.pdf', bbox_inches='tight')
# plt.savefig('fresnel_cross.png', bbox_inches='tight')
plt.show()
# Closing down the script and stopping the timer.
end = time.time()
print('The program took {:.3f} seconds!'.format((end-start)))
test cases (standalone)/fourierOptics_03_verification.py
+775
View File
@@ -0,0 +1,775 @@
#!/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
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, _y, _E_out.real, kind='linear')
_f_ip_imag = interpolate.interp2d(_x, _y, _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
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 phaseDelayFunction(self, x, y, l, paraxial=False):
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
if paraxial:
_dxy = self.thickness - (x**2 + y**2)/2 * (1/self.R1 - 1/self.R2)
_k = 2 * np.pi / 30e-3
_n = self.findRefractiveIndex(self.referenceWavelength)
_t_l = np.exp(-1j * _k * self.thickness)* np.exp(-1j * _k * (_n - 1) * _dxy)
# 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 phaseDelayFunction(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 phaseDelayFunction(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 phaseDelayFunction(self, l):
_k = 2 * np.pi / l
_t_m = np.exp(-1j * _k * 2 * self.z)
return _t_m
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 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)
_wf = U_in
for _lens in self.opticalTrain:
_i = self.opticalTrain.index(_lens)
_dz = self.sequentialDistance[_i]
_wf = _lens.phaseDelayFunction(x, y, l, paraxial=paraxial)
# Perform Fresnel number test:
_F = self.fresnel_number()
if _F < 5:
# Fresnel propagation
_wf, x, y = self.fresnel_propagation()
elif _F >= 5:
# Angular spectrum (p2p) propagation
_wf, x, y = self.angular_spectrum_propagation()
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])
def FGB_1D(w0,l,z,x):
k = 2 * np.pi / l
z = z + 1e-16
w = w0 * np.sqrt( 1 + ( (l*z) / (np.pi*w0**2) )**2)
R = z + 1/z * ( np.pi * w0**2 / l)**2
phi = np.arctan2( (l * z) , (np.pi * w0**2) )
norm = (2./ np.pi)**(.25) * (1. / w0**2)**(.5)
beam = np.exp( -x**2/w**2 - (1j * k * z) - ((1j*np.pi*x**2)/(l*R)) + 1j*phi/2)
return beam * norm
if __name__ == '__main__':
plt.rc('font', family='serif')
plt.rc('axes', labelsize='large')
import time
start = time.time()
l = 30e-3
# l = 1e-6
fou = FourierOptics(l=l)
dpix = l/6
n = int(500 / dpix * l)
m = int(500 / dpix * l)
n, m = 3001, 3001
print (n,m)
prof = profiles()
# A = prof.topHat(4*l, 2001, 2001, 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))
w0_init = 4*l
z = 400*l
wg = w0_init * np.sqrt( 1 + ( (l*z) / (np.pi*w0_init**2) )**2)
# rmax_init = 4 * w0_init
# rmax_init = 1000/8*l
xGauss = np.linspace(-(n/2)*dpix, (n/2)*dpix, n)
rmax_init = np.max(xGauss)
# A = GaussianBeam(w0_init, l).fundamentalMode(z, n, m, rmax_init)
A0 = FGB_1D(w0_init, l, 0, xGauss)
A0 /= np.sqrt(np.sum(np.abs(A0)**2))
A = FGB_1D(wg, l, z, xGauss)
A /= np.sqrt(np.sum(np.abs(A)**2))
A2D = GaussianBeam(w0_init, l).fundamentalMode(z, n, m, rmax_init)
A2D /= np.sqrt(np.sum(np.abs(A2D[int(n/2)+1,:])**2))
A2D0 = GaussianBeam(w0_init, l).fundamentalMode(0, n, m, rmax_init)
Ap2p,xp2p,_ = fou.planeWavePropagationFFTW(A2D0, dz=z, dpix=dpix, show=False, save=False)
Ap2p /= np.sqrt(np.sum(np.abs(Ap2p[int(n/2),:])**2))
oC = [xGauss.min(), xGauss.max(), xGauss.min(), xGauss.max()]
# savename= 'FourierOptics_03_verification_fresnel.pdf'
savename=None
Afren,xfren,_ = fou.FresnelPropagationFFTW(A2D0, x,y, dz=z, show=True, save=savename, outputCoordinates=oC)
Afren /= np.sqrt(np.sum(np.abs(Afren[int(n/2),:])**2))
# Afren = np.conj(Afren)
norm_fac = np.max(np.abs(A2D)[int(n/2),:])
fig = plt.figure()
ax = fig.add_subplot(111)
# ax.plot(np.abs(A0), label='z=0, analytical')
# ax.plot(np.abs(A)/norm_fac, label='z=%i, analytical' %(int(z/l)))
ax.plot(np.abs(A2D)[int(n/2),:]/norm_fac, 'r--', label='z=%i$\lambda$, analytical' %(int(z/l)))
ax.plot(np.abs(Ap2p)[int(n/2),:]/norm_fac, label='z=%i$\lambda$, plane to plane' %(int(z/l)))
ax.plot(np.abs(Afren)[int(n/2),:]/norm_fac, label='z=%i$\lambda$, Fresnel' %(int(z/l)))
ax.grid(alpha=0.3)
ax.set_xlabel('Location [pixel no.]')
ax.set_ylabel('Normalized amplitude')
ax.legend()
plt.savefig('FourierOptics_03_verification_amplitude.pdf', bbox_inches='tight', transparent=True)
fig2 = plt.figure(figsize=(12,3))
# ax1 = fig2.add_subplot(151)
# ax1.plot(np.angle(A0), label='z=0, analytical')
# ax1.set_title('z=0, analytical')
# ax2 = fig2.add_subplot(131)
# ax2.plot(np.angle(A), label='z=%i, analytical' %(int(z/l)))
# ax2.set_title('z=%i, analytical' %(int(z/l)))
ax3 = fig2.add_subplot(131)
ax3.plot(np.angle(A2D)[int(n/2),:], label='z=%i$\lambda$, analytical' %(int(z/l)))
ax3.set_title('z=%i$\lambda$, Gaussian beam' %(int(z/l)))
ax3.set_ylabel('Phase')
ax3.set_xlabel('Location [pixel no.]')
ax3.grid(True, alpha=0.3)
ax4 = fig2.add_subplot(132)
ax4.plot(np.angle(Ap2p)[int(n/2),:], label='z=%i$\lambda$, plane to plane' %(int(z/l)))
ax4.set_title('z=%i$\lambda$, Plane to plane' %(int(z/l)))
ax4.set_ylabel('Phase')
ax4.set_xlabel('Location [pixel no.]')
ax4.grid(True, alpha=0.3)
ax5 = fig2.add_subplot(133)
ax5.plot(np.angle(Afren)[int(n/2),:], label='z=%i$\lambda$, Fresnel' %(int(z/l)))
ax5.set_title('z=%i$\lambda$, Fresnel' %(int(z/l)))
ax5.set_ylabel('Phase')
ax5.set_xlabel('Location [pixel no.]')
ax5.grid(True, alpha=0.3)
# ax1.grid(alpha=0.3)
# ax6 = fig2.add_subplot(166)
# ax6.legend(ax=ax6)
plt.savefig('FourierOptics_03_verification_phase.pdf', bbox_inches='tight', transparent=True)
plt.tight_layout()
plt.show()
# set_trace()
# Afren,_,_ = fou.FresnelPropagationFFTW(A2D, dz=400*l, dpix=dpix, oversample=1, show=True, save=False, outputCoordinates=oC)
gauss = GaussianBeam(w0_init, l)
# A = prof.topHat(4*l, 2001, 2001, dpix=dpix, oversample=1)
A = gauss.fundamentalMode(0,n,m,rmax_init)
if False:
Ics = []
Ics_x = []
Isum = []
lList = np.linspace(1,201,10).tolist()
for i in lList:
title = f"{lList.index(i):03d}"
print (' Calculating {0}/{1}'.format(title, f"{len(lList):03d}"), end="\r")
# Exyz, x, y = fou.planeWavePropagationFFTW(A, dz=i*l, dpix=dpix, show=False, save=False)#, scaleInputArray=(n,m))
Exyz, x, y = fou.FresnelPropagationFFTW(A, dz=i*l, dpix=dpix, show=False, save=False, outputCoordinates=oC)
# Exyz = gauss.fundamentalMode(i*l, n, m, rmax_init)
# y, x = np.mgrid[-rmax_init:rmax_init:n*1j,-rmax_init:rmax_init:n*1j]
Isum.append(np.sum(np.abs(Exyz)**2))
Ics.append((np.abs(Exyz[int(Exyz.shape[0]/2)])**2))
Ics_x.append(x.real[0])
y_ext = [y.real[0,0],y.real[-1,0]]
print(y_ext)
Exyz = None
x = None
y = None
fig = plt.figure(figsize=(8,4))
ax = fig.add_subplot(211)
ax.tick_params(axis ='both', direction='in')
im = ax.imshow(np.asarray(Ics).T, cmap='hot', aspect='auto', interpolation='bicubic', extent=[lList[0],lList[-1],y_ext[0],y_ext[1]])
ax.yaxis.set_major_formatter(FormatStrFormatter('%.1E'))
# ax.set_ylim([-2.5e-4,2.5e-4])
ax2 = fig.add_subplot(212)
ax2.plot(lList, np.asarray(Ics)[:,int(np.asarray(Ics).shape[1]/2)], 'k-')
ax2.yaxis.set_major_formatter(FormatStrFormatter('%.1E'))
ax2.set_xlabel('Propagation distance [$\lambda$]')
ax2.tick_params(axis ='both', direction='in')
ax2.set_ylabel('Intensity')
ax2.set_xlim([np.min(lList), np.max(lList)])
ax2.grid(True, alpha=0.3)
# ax3 = fig.add_subplot(133)
# ax3.set_ylabel('Integrated intensity')
# ax3.set_xlabel('Propagation distance [$\lambda$]')
# ax3.grid(True, alpha=0.3)
# ax3.tick_params(axis ='both', direction='in')
# ax3.plot(Isum)
plt.tight_layout()
# plt.savefig('fresnel_cross.pdf', bbox_inches='tight')
# plt.savefig('fresnel_cross.png', bbox_inches='tight')
plt.show()
# E_xyz,_,_ = fou.planeWavePropagationFFTW(A, dz=4*l, dpix=dpix, oversample=1, show=True, save=False, outputLimits=(-0.4*1e-3,0.4*1e-3))#, scaleInputArray=(n,m))
# Exyz,_,_ = fou.FresnelPropagationFFTW(A, dz=4*l, dpix=dpix, oversample=1, show=True, save=False)#, scaleInputArray=(n,m))
end = time.time()
print((end-start)*1000)
# set_trace()
test cases (standalone)/fourierOptics_04_simple_lens_system.py
+720
View File
@@ -0,0 +1,720 @@
#!/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)))
test cases (standalone)/fourierOptics_25_single_lens_system.py
+729
View File
@@ -0,0 +1,729 @@
#!/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, _y, _E_out.real, kind='linear')
_f_ip_imag = interpolate.interp2d(_x, _y, _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
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 phaseDelayFunction(self, x, y, l, paraxial=False):
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
if paraxial:
_dxy = self.thickness - (x**2 + y**2)/2 * (1/self.R1 - 1/self.R2)
_k = 2 * np.pi / 30e-3
_n = self.findRefractiveIndex(self.referenceWavelength)
_t_l = np.exp(-1j * _k * self.thickness)* np.exp(-1j * _k * (_n - 1) * _dxy)
# 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 phaseDelayFunction(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 phaseDelayFunction(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 phaseDelayFunction(self, l):
_k = 2 * np.pi / l
_t_m = np.exp(-1j * _k * 2 * self.z)
return _t_m
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 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)
_wf = U_in
for _lens in self.opticalTrain:
_i = self.opticalTrain.index(_lens)
_dz = self.sequentialDistance[_i]
_wf = _lens.phaseDelayFunction(x, y, l, paraxial=paraxial)
# Perform Fresnel number test:
_F = self.fresnel_number()
if _F < 5:
# Fresnel propagation
_wf, x, y = self.fresnel_propagation()
elif _F >= 5:
# Angular spectrum (p2p) propagation
_wf, x, y = self.angular_spectrum_propagation()
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/50
n, m = 2001, 2001
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 = profiles()
# A = prof.spiderTopHat(250, 100, 2001, 2001, dpix=1, oversample=1)
A = prof.topHat(2.5, n, m, dpix=dpix, oversample=1)
# Alternatively, create a gaussian profile
# gauss = GaussianBeam(0.3, l*1000)
# A = gauss.fundamentalMode(0,n,m,np.max(x))
# 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 a focal length ~100 mm, with a clear aperture radius of 4 mm and thickness of 1 mm.
L1 = lensObject('N-BK7', 50, 1e99, 1, 4, name='L1')
print('Focal length of lens {} is {:.1f} mm.'.format(L1.name, L1.f))
# Calculate the phase delay introduced by the lens
t_l = L1.phaseDelayFunction(x,y,1.35e-6, paraxial=False)
# The input distribution gets a curved wavefront after the lens
U_in = A * t_l
# Performing the propagation using Fresnel approximation to the focal plane.
oC = [-5, 5, -5, 5]
savename = 'simple_lens.pdf'
Exyz,_,_ = fou.FresnelPropagationFFTW(U_in, x, y, dz=L1.f, oversample=None, show=True, save=savename, outputCoordinates=oC)
# Plot a cross section of the new field (useful to better see the diffraction rings)
fig = plt.figure()
ax = fig.add_subplot(111)
ax.grid(True, alpha=0.3)
ax.semilogy(np.linspace(np.min(oC), np.max(oC),n), np.abs(Exyz[1000,:]**2), 'r-')
ax.set_xlabel('x (mm)')
ax.set_ylabel('Intensity')
ax.tick_params(axis ='both', direction='in', which='both')
plt.savefig('simple_lens_slice.pdf', bbox_inches='tight', transparent=True)
plt.show()
# The following produces a cross-section of the propagation for various z, so that you can see how the beam evolves.
if True:
Ics = []
Ics_x = []
Isum = []
xmax = []
zList = np.linspace(20,220,20).tolist()
xin, yin = x, y
for i in zList:
title = f"{zList.index(i)+1:03d}"
print (' Calculating {0}/{1}'.format(title, f"{len(zList):03d}"), end="\r")
# Propagate the wavefront
# Exyz, x, y = fou.planeWavePropagationFFTW(A, dz=i*l, dpix=dpix, oversample=1, show=False, save=False)#, scaleInputArray=(n,m))
Exyz, x, y = fou.FresnelPropagationFFTW(U_in, xin, yin, dz=i, oversample=None, show=False, save=False, outputCoordinates=oC)
# Process the results by taking a slice of the data
Isum.append(np.sum(np.abs(Exyz)**2))
Ics.append((np.abs(Exyz[int(Exyz.shape[0]/2)])**2))
Ics_x.append(x.real[0])
y_ext = [y.real[0,0],y.real[-1,0]]
# Reset the data
x = None
y = None
# Plotting all slices
fig = plt.figure(figsize=(8,4))
ax = fig.add_subplot(211)
ax.tick_params(axis ='both', direction='in')
# Plotting the square root of the amplitude for better contrast
im = ax.imshow(np.sqrt(np.asarray(Ics)).T, cmap='hot', aspect='auto', interpolation='bicubic', extent=[zList[0],zList[-1],y_ext[0],y_ext[1]])
ax.yaxis.set_major_formatter(FormatStrFormatter('%.1E'))
# Plotting the intensity relative to the input
ax2 = fig.add_subplot(212)
ax2.plot(zList, np.asarray(Ics)[:,int(np.asarray(Ics).shape[1]/2)], 'k-')
ax2.yaxis.set_major_formatter(FormatStrFormatter('%.1E'))
ax2.set_xlabel('Propagation distance (mm)')
ax2.tick_params(axis ='both', direction='in')
ax2.set_ylabel('Amplitude')
ax2.set_xlim([np.min(zList), np.max(zList)])
ax2.grid(True, alpha=0.3)
# Saving the figure? Then show.
plt.tight_layout()
# plt.savefig('fresnel_cross.pdf', bbox_inches='tight')
# plt.savefig('fresnel_cross.png', bbox_inches='tight')
plt.savefig('simple_lens_cross.pdf', bbox_inches='tight', transparent=True)
plt.show()
# Closing down the script and stopping the timer.
end = time.time()
print('The program took {:.3f} seconds!'.format((end-start)))
test cases (standalone)/glassCatalog.py
+12
View File
@@ -0,0 +1,12 @@
#!/usr/bin/env Python3
# Pyton dictionary with some glass types
catalog = {
'N-BK7' : [1.039612120E+00, 6.000698670E-03, 2.317923440E-01, 2.001791440E-02, 1.010469450E+00, 1.035606530E+02 ],
'SiO2' : [0.6961663, 0.0684043, 0.4079426, 0.1162414, 0.8974794, 9.896161 ],
'F2' : [1.345333590E+00, 9.977438710E-03, 2.090731760E-01, 4.704507670E-02, 9.373571620E-01, 1.118867640E+02 ],
'S-FPL51' : [1.170105050E+00, 6.162039240E-03, 4.757107830E-02, 2.633728760E-02, 7.638324450E-01, 1.418826420E+02 ],
'S-LAH71' : [1.982800310E+00, 1.189874560E-02, 3.167584500E-01, 5.271560010E-02, 2.444726460E+00, 2.132206970E+02 ],
}