mirror of
https://github.com/yosata/fourierOptics.git
synced 2026-10-30 05:01:04 +00:00
227 lines
8.4 KiB
Python
227 lines
8.4 KiB
Python
#!/usr/bin/env Python3
|
|
import numpy as np
|
|
from pdb import set_trace
|
|
import json
|
|
import matplotlib.pyplot as plt
|
|
from matplotlib.ticker import FormatStrFormatter
|
|
from matplotlib.colors import LogNorm
|
|
import sys
|
|
sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics')
|
|
from fourierOptics import *
|
|
|
|
def saveVariableToTXT(variable, filename):
|
|
variableFile = open(filename, 'w')
|
|
|
|
if type(variable) == np.ndarray:
|
|
json.dump(variable.tolist(), variableFile)
|
|
elif type(variable) in [list, str, float]:
|
|
json.dump(variable, variableFile)
|
|
else:
|
|
print ('Not a list, array or string. Check variable type.')
|
|
|
|
variableFile.close()
|
|
return
|
|
|
|
# Load a variable from a json text file.
|
|
def loadVariableFromTXT(filename):
|
|
variableFile = open(filename, 'r')
|
|
variable = json.load(variableFile)
|
|
variableFile.close()
|
|
return variable
|
|
|
|
def scale_zernike_coefficients(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))
|
|
return c
|
|
|
|
def expected_strehl(rms, l):
|
|
return np.exp( -(2*np.pi*rms/l)**2 )
|
|
|
|
if __name__ == '__main__':
|
|
plt.rc('font', family='serif')
|
|
plt.rc('axes', labelsize='large')
|
|
|
|
# Parameters
|
|
dpix = 160
|
|
n, m = 1001, 1001
|
|
oC = [-.5, .5, -.5, .5]
|
|
l = 1.35e-3 # wavelength
|
|
zern_rms = 45e-5 # RMS of the zernike component of the WFE
|
|
midspat_rms = 21.8e-5 # RMS of the PSD component of the WFE
|
|
rmsList = np.linspace(1e-4, 1e-3, 25) # 1 to 1000 nm WFE RMS in mm
|
|
number_of_iterations = 50
|
|
|
|
# Start of the ELT-lens system
|
|
# Initiate Fourier Optics, profile and wavefront error objects
|
|
fou = FourierOptics(l=l)
|
|
prof = Profiles()
|
|
wfe = WavefrontError(n,m)
|
|
|
|
print ('Your resolution is {}x{} with a pixel size of {}'.format(n,m,dpix))
|
|
A = prof.top_hat(38854/2, n, m, dpix=dpix, oversample=1)
|
|
|
|
# Create an input grid, equivalent to the size of ELT-M1
|
|
x, y = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m))
|
|
|
|
# Create a ELT-M1 parabolic mirror
|
|
M1 = ParabolicMirror(dx=0, dy=0, name='ELT-M1', CA=38854/2, f=679945)
|
|
t_l = M1.phase_delay(x,y,l)
|
|
U_in = A * t_l
|
|
|
|
# Propagating the original, unperturbed, system and calculating encircled energy
|
|
Exyz_og,xp,yp = fou.fresnel_propagation(U_in, x, y, dz=M1.f, show=False, outputCoordinates=oC)
|
|
|
|
# Initiate psf fitter object to calculate the strehl ratios
|
|
psf = PSFFitter(xp, yp, Exyz_og)
|
|
|
|
StrehlList_zern_complete, StrehlList_psd_complete = [], []
|
|
rms_zern_list_complete, rms_psd_list_complete = [], []
|
|
|
|
for i in range(number_of_iterations):
|
|
print('Iteration {}/{}'.format(i+1,number_of_iterations))
|
|
|
|
# Generate a PSD surface
|
|
psd_surface = wfe.generate_height_map(midspat_rms, 2, 1)
|
|
|
|
# Fit the PSD surface using zernike terms
|
|
# Generating a normalized grid for the Zernike evaluation
|
|
yz,xz = np.mgrid[-1:1:wfe.n*1j, -1:1:wfe.m*1j]
|
|
_r = np.sqrt(xz**2 + yz**2)
|
|
_mask = _r > 1
|
|
coeffic = wfe.fit_zernike_coefficients(psd_surface, modes=66)
|
|
|
|
StrehlList_zern, StrehlList_psd = [], []
|
|
rms_zern_list, rms_psd_list = [], []
|
|
|
|
for rms in rmsList:
|
|
# Make a new perturbed wavefront error with the required RMS
|
|
coeffic = wfe.scale_zernike_coefficients(coeffic, rms)
|
|
hzern = wfe.zernike_cartesian(coeffic, xz, yz)
|
|
|
|
# Add the new wavefront error
|
|
M1.add_wavefront_error(hzern)
|
|
t_l = M1.phase_delay(x,y,l)
|
|
U_in = A * t_l
|
|
|
|
# Performing the propagation using Fresnel approximation to the focal plane.
|
|
Exyz_zern,xp,yp = fou.fresnel_propagation(U_in, x, y, dz=M1.f, show=False, outputCoordinates=oC)
|
|
Strehl_zern = np.abs(psf.get_coupling(Exyz_og, Exyz_zern))**2
|
|
|
|
set_trace()
|
|
StrehlList_zern.append(Strehl_zern)
|
|
|
|
####################################
|
|
# Add the new wavefront error
|
|
factor = rms / midspat_rms
|
|
h_psd = wfe.scale_rms(factor, data=psd_surface)
|
|
|
|
# Add the new wavefront error
|
|
M1.add_wavefront_error(h_psd)
|
|
t_l = M1.phase_delay(x,y,l)
|
|
U_in = A * t_l
|
|
|
|
# Performing the propagation using Fresnel approximation to the focal plane.
|
|
Exyz_psd,xp,yp = fou.fresnel_propagation(U_in, x, y, dz=M1.f, show=False, outputCoordinates=oC)
|
|
Strehl_psd = np.abs(psf.get_coupling(Exyz_og, Exyz_psd))**2
|
|
StrehlList_psd.append(Strehl_psd)
|
|
|
|
rms_zern = np.sqrt(np.mean(np.asarray(hzern)**2))
|
|
rms_psd = np.sqrt(np.mean(np.asarray(h_psd)**2))
|
|
print('RMS: {:.3f}, S_zern: {:.3f}, S_psd: {:.3f}, \
|
|
RMS_zern: {:.3f}, RMS_psd: {:.3f}'.format(
|
|
rms*1e6, Strehl_zern, Strehl_psd, rms_zern*1e6, rms_psd*1e6))
|
|
rms_zern_list.append(rms_zern)
|
|
rms_psd_list.append(rms_psd)
|
|
|
|
StrehlList_zern_complete.append(StrehlList_zern)
|
|
StrehlList_psd_complete.append(StrehlList_psd)
|
|
rms_zern_list_complete.append(rms_zern_list)
|
|
rms_psd_list_complete.append(rms_psd_list)
|
|
|
|
|
|
mean_strehl_zern = np.mean(StrehlList_zern_complete, axis=0)
|
|
err_strehl_zern = np.std(StrehlList_zern_complete, axis=0)
|
|
mean_strehl_psd = np.mean(StrehlList_psd_complete, axis=0)
|
|
err_strehl_psd = np.std(StrehlList_psd_complete, axis=0)
|
|
|
|
rms_zern_list_complete = np.asarray(rms_zern_list_complete)
|
|
rms_psd_list_complete = np.asarray(rms_psd_list_complete)
|
|
|
|
# saveVariableToTXT(rmsList, 'rmsList.txt')
|
|
# saveVariableToTXT(mean_strehl_zern, 'mean_strehl_zern.txt')
|
|
# saveVariableToTXT(err_strehl_zern, 'err_strehl_zern.txt')
|
|
# saveVariableToTXT(mean_strehl_psd, 'mean_strehl_psd.txt')
|
|
# saveVariableToTXT(err_strehl_psd, 'err_strehl_psd.txt')
|
|
# saveVariableToTXT(rms_psd_list_complete, 'rms_psd_list_complete.txt')
|
|
# saveVariableToTXT(rms_zern_list_complete, 'rms_zern_list_complete.txt')
|
|
|
|
|
|
# Normal plot
|
|
fig = plt.figure()
|
|
ax = fig.add_subplot(111)
|
|
for i in range(number_of_iterations):
|
|
# ax.plot(rmsList*1e6, StrehlList_zern_complete[i], 'y--', alpha=0.2)
|
|
# ax.plot(rmsList*1e6, StrehlList_psd_complete[i], 'b--', alpha=0.2)
|
|
ax.plot(rms_zern_list_complete[i]*1e6, StrehlList_zern_complete[i], 'y-', alpha=0.1)
|
|
ax.plot(rms_psd_list_complete[i]*1e6, StrehlList_psd_complete[i], 'b-', alpha=0.1)
|
|
ax.plot(rmsList*1e6, expected_strehl(rmsList,l), 'k-', label='Theory')
|
|
# ax.errorbar(rmsList*1e6, mean_strehl_zern, yerr=err_strehl_zern, color='y', linestyle='-', label='Zernikes', capsize=3, markeredgewidth=1)
|
|
# ax.errorbar(rmsList*1e6, mean_strehl_psd, yerr=err_strehl_psd, color='b', linestyle='-', label='PSD', capsize=3, markeredgewidth=1)
|
|
ax.errorbar(rmsList*1e6, mean_strehl_zern, yerr=err_strehl_zern, color='y', linestyle='-', label='Zernikes', capsize=3, markeredgewidth=1)
|
|
ax.errorbar(rmsList*1e6, mean_strehl_psd, yerr=err_strehl_psd, color='b', linestyle='-', label='PSD', capsize=3, markeredgewidth=1)
|
|
ax.set_ylim([0,1])
|
|
ax.set_xlim([np.min(rmsList)*1e6, np.max(rmsList)*1e6])
|
|
ax.grid(True, alpha=0.3)
|
|
ax.set_xlabel('RMS (nm)')
|
|
ax.set_ylabel('Strehl')
|
|
ax.legend()
|
|
plt.tight_layout()
|
|
# plt.savefig('strehl_verification.pdf', bbox_inches='tight', transparent=True)
|
|
|
|
|
|
|
|
|
|
# Log plot
|
|
slSquared = (rmsList / l)**2
|
|
err_strehl_zern = 0.434 * err_strehl_zern / mean_strehl_zern
|
|
err_strehl_psd = 0.434 * err_strehl_psd / mean_strehl_psd
|
|
|
|
|
|
|
|
fig = plt.figure()
|
|
ax = fig.add_subplot(111)
|
|
for i in range(number_of_iterations):
|
|
ax.plot( (rms_zern_list_complete[i]/l)**2, np.log10(StrehlList_zern_complete[i]), 'y-', alpha=0.1)
|
|
ax.plot( (rms_psd_list_complete[i]/l)**2, np.log10(StrehlList_psd_complete[i]), 'b-', alpha=0.1)
|
|
ax.plot(slSquared, np.log10(expected_strehl(rmsList,l)), 'k-', label='Theory')
|
|
ax.errorbar(slSquared, np.log10(mean_strehl_zern), yerr=err_strehl_zern, color='y', linestyle='-', label='Zernikes', capsize=3, markeredgewidth=1)
|
|
ax.errorbar(slSquared, np.log10(mean_strehl_psd), yerr=err_strehl_psd, color='b', linestyle='-', label='PSD', capsize=3, markeredgewidth=1)
|
|
# ax.set_ylim([0,1])
|
|
ax.set_xlim([np.min(slSquared), np.max(slSquared)])
|
|
ax.grid(True, alpha=0.3)
|
|
ax.set_xlabel('$\\left(\\sigma/\\lambda\\right)^2$')
|
|
ax.set_ylabel('$\\log_{10}(S)$')
|
|
ax.legend()
|
|
plt.tight_layout()
|
|
# plt.savefig('strehl_verification_log.pdf', bbox_inches='tight', transparent=True)
|
|
|
|
|
|
|
|
# Show the last psd and zernike surface that were generated
|
|
psd_surface[_mask] = np.nan
|
|
hzern[_mask] = np.nan
|
|
fig2 = plt.figure(figsize=(10,4))
|
|
ax1 = fig2.add_subplot(121)
|
|
ax2 = fig2.add_subplot(122)
|
|
im1 = ax1.imshow(hzern, cmap='jet')
|
|
im2 = ax2.imshow(psd_surface, cmap='jet')
|
|
plt.tight_layout()
|
|
# plt.savefig('wavefront_error_example_with_zern_fit.pdf', bbox_inches='tight', transparent=True)
|
|
plt.show()
|
|
|
|
# set_trace()
|
|
|