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
Files
master
fourierOptics/scripts/MICADO_tilt_fit.py
T

409 lines
16 KiB
Python
Raw Permalink Blame History

#!/usr/bin/env Python3
import numpy as np
from pdb import set_trace
import matplotlib.pyplot as plt
from matplotlib.ticker import FormatStrFormatter
from matplotlib.colors import LogNorm
import time
import sys
sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics')
from fourierOptics import *
from scipy import interpolate
'''
This file simulates the paraxial optical model of the ELT-MICADO system
with an ADC placed inside the collimated beam near the pupil.
Surface form errors are added to the prisms to investigate the effect of
WFEs on the PSF centroid.
Author: Joost van den Born
Contact: born@astron.nl
'''
class ADC_surfaces:
'''
Very specific class to generate and scale wavefront errors on
the MICADO ADC prisms.
'''
def __init__(self, x,y, clearAperture=125):
self.n, self.m = x.shape
self.clearAperture = clearAperture
self.radius = self.clearAperture / 2
self.yz,self.xz = np.mgrid[x.min()/self.radius:x.max()/self.radius:self.n*1j,
x.min()/self.radius:x.max()/self.radius:self.m*1j]
self.zern_coefficients = None
self.zern_surface = None
self.psd_surface = None
self.latest_surface = None
def generate_ADC_wfe(self, zernRMS=20e-6, psdRMS=15e-6, zernModes=19):
'''
We apply the wavefront error an each of the ADC prisms
as a phase transformation, corresponding to a
transmitted wavefront error.
This wfe is separated into a zernike and psd component
with a total rms of 25 nm
zernike component - RMS: 20 nm
psd component - RMS: 15 nm
Total - RMS: 25 nm
'''
print('Generating new ADC WFE...')
wfe = WavefrontError(self.n,self.m)
zern_rms = zernRMS # RMS of the zernike component of the WFE
midspat_rms = psdRMS # RMS of the PSD component of the WFE
self.psd_surface = wfe.generate_height_map(midspat_rms, 2, 1)
zern_surface = wfe.generate_height_map(zern_rms, 2, 1)
zern_coefficients = wfe.fit_zernike_coefficients(zern_surface, modes=zernModes)
zern_coefficients[:3] = 0 # No defocus, tip or tilt
self.zern_coefficients = wfe.scale_zernike_coefficients(zern_coefficients, zern_rms)
# Scale the zernike modes to the clear aperture of the ADC
self.zern_surface = wfe.zernike_cartesian(self.zern_coefficients, self.xz, self.yz, rescale=False)
zern_surface = self.mask_surface(self.zern_surface)
psd_surface = self.mask_surface(self.psd_surface)
self.latest_surface = zern_surface + psd_surface
return self.latest_surface
def mask_surface(self, A):
'''
Set the array values to 0 where
the radius is > 1, without changing
the original array.
'''
# Rescale x and y
_mask = np.sqrt(self.xz**2 + self.yz**2) > 1
B = A.copy()
B[_mask] = 0
return B
def scale_wfe(self, zernRMS=20e-5, psdRMS=15e-5):
'''
Scaling of the original wfe, given the new
zernike and psd RMS values that are desired.
'''
print('Scaling ADC WFE...')
_old_psd_rms = np.sqrt(np.mean(self.psd_surface**2))
_factor_psd = psdRMS / _old_psd_rms
wfe = WavefrontError(self.n,self.m)
new_psd_surface = wfe.scale_rms(_factor_psd, data=self.psd_surface)
new_zern_coefficients = wfe.scale_zernike_coefficients(self.zern_coefficients, zernRMS)
# Scale the zernike modes to the clear aperture of the ADC
new_zern_surface = wfe.zernike_cartesian(new_zern_coefficients, self.xz, self.yz, rescale=False)
new_zern_surface = self.mask_surface(new_zern_surface)#, _xz, _yz)
new_psd_surface = self.mask_surface(new_psd_surface)#, _xz, _yz)
self.latest_surface = new_zern_surface + new_psd_surface
return self.latest_surface
def fit_wfe_tilt(self,z=None):
'''
Takes the WFE on the ADC, finds the clear aperture
and fits only the tilt over that area.
'''
if z is None:
z = self.latest_surface
_, _ix1 = find_nearest(self.xz,-1)
_, _ix2 = find_nearest(self.xz, 1)
_zselect = z[_ix1:_ix2,_ix1:_ix2]
_wfe = WavefrontError(_zselect.shape[0], _zselect.shape[1])
_C = _wfe.fit_zernike_coefficients(_zselect, modes=3)
_fit_surface = _wfe.zernike_cartesian(_C, self.xz, self.yz, rescale=False)
_fit_surface = self.mask_surface(_fit_surface)
return _fit_surface, _C
if __name__ == '__main__':
plt.rc('font', family='serif')
plt.rc('axes', labelsize='large')
l = 1.35e-3
fou = FourierOptics(l=l)
# MICADO paraxial system
# Pick sampling parameters
dpix = 50
n, m = 3001, 3001
print ('Your resolution is {}x{} with a pixel size of {}'.format(n,m,dpix))
# 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))
print('This means your grid is {} mm by {} mm, with each pixel {} mm.'.format(x.max()-x.min(), y.max()-y.min(), dpix))
fList = [1]#[0.001, 0.01, 0.1, 0.5, 1, 2, 3, 5, 7, 10]
###########################################################
#################### SETTING UP MICADO ####################
###########################################################
# Define the ELT mask (the spiders and central obscuration)
prof = Profiles()
# ELT_MASK = prof.top_hat(19427, n,m,dpix)#.astype('float64')
# ELT_MASK = prof.obscured_top_hat(4000, 500, n, m, dpix)
# ELT_MASK = prof.ELT_pupil(n,m,dpix)
# ELT Mirror system combined into a single lens
ELT_MIRROR = ParabolicMirror(dx=0, dy=0, name='ELT-paraxial', CA=38854, f=6.84022e5)
ELT_FP = OpticalSurface(name='MCD-IFP')
ELT_COL = ParaxialLens(CA=154.54, f=1418.94, name='MCD-COL')
ELT_ADC = ADCObject()
ELT_CAM = ParaxialLens(CA=89.2, f=4278.32701, name='MCD-CAM')
ELT_DET = OpticalSurface(name='MCD-DET')
MCD = OpticalTrain()
MCD.add_object(ELT_MIRROR, dist=0)
MCD.add_object(ELT_FP, dist=ELT_MIRROR.f)
MCD.add_object(ELT_COL, dist=ELT_COL.f)
ELT_ADC.add_to_optical_train(MCD, 1209.4718)
MCD.add_object(ELT_CAM, dist=50)
MCD.add_object(ELT_DET, dist=ELT_CAM.f)
# Print lens data for verification
MCD.print_lens_data()
###########################################################
# Load propagation up to the ADC (so we don't need to compute this every time)
Exyz3, x3, y3 = load_variable('./data/ADC_input_(1.35).pkl')
x3 = x3.real
y3 = y3.real
E0, x0, y0 = load_variable('./data/MICADO_PSF_(1.35).pkl')
###########################################################
#################### ADD WFE TO ADC #######################
###########################################################
# adc_wfe = ADC_surfaces(x3,y3,125)
# wfe_p1 = adc_wfe.generate_ADC_wfe()
# newsurf = adc_wfe.scale_wfe()
wfe_p1 = ADC_surfaces(x3,y3)
wfe_p2 = ADC_surfaces(x3,y3)
wfe_p3 = ADC_surfaces(x3,y3)
wfe_p4 = ADC_surfaces(x3,y3)
# wfe_p2 = generate_ADC_wfe(x3, y3)
# wfe_p3 = generate_ADC_wfe(x3, y3)
# wfe_p4 = generate_ADC_wfe(x3, y3)
wfe_p1.generate_ADC_wfe()
wfe_p2.generate_ADC_wfe()
wfe_p3.generate_ADC_wfe()
wfe_p4.generate_ADC_wfe()
fit1, _ = wfe_p1.fit_wfe_tilt()
fit2, _ = wfe_p2.fit_wfe_tilt()
fit3, _ = wfe_p3.fit_wfe_tilt()
fit4, _ = wfe_p4.fit_wfe_tilt()
# Do the analysis for different scaling factors
popt_airy_list, popt_gauss_list, popt_moffat_list, popt_cust_list = [], [], [], []
pcov_airy_list, pcov_gauss_list, pcov_moffat_list, custom_var_list = [], [], [], []
strehl_list = []
for f in [0,1]:
print('Factor set to {:.1f}'.format(f))
ELT_ADC.p1.add_wavefront_error([wfe_p1.latest_surface, fit1][f])
ELT_ADC.p2.add_wavefront_error([wfe_p2.latest_surface, fit2][f])
ELT_ADC.p3.add_wavefront_error([wfe_p3.latest_surface, fit3][f])
ELT_ADC.p4.add_wavefront_error([wfe_p4.latest_surface, fit4][f])
###########################################################
###########################################################
###################### PROPAGATION ########################
###########################################################
coordinatesList = [
[-1.5, 1.5, -1.5, 1.5],
[-200, 200, -200, 200],
[-120, 120, -120, 120],
[-120, 120, -120, 120],
[-120, 120, -120, 120],
[-120, 120, -120, 120],
[-120, 120, -120, 120],
[-2.5, 2.5, -2.5, 2.5]
]
print('Starting propagation...')
# From MCD-ADC-P1 to MCD-ADC-P2
print('Propagating from MCD-ADC-P1 to MCD-ADC-P2')
Exyz3_ = Exyz3 * ELT_ADC.get_phase_delay(x3,y3,l, theta=None, paraxial=True, prism=1)
Exyz4, x4, y4 = fou.plane_wave_propagation(Exyz3_, x3, y3, dz=MCD.sequentialDistance[4], oversample=None,
show=False, save=False, outputCoordinates=None)
# From MCD-ADC-P2 to MCD-ADC-P3
print('Propagating from MCD-ADC-P2 to MCD-ADC-P3')
Exyz4 *= ELT_ADC.get_phase_delay(x4,y4,l, theta=None, paraxial=True, prism=2)
Exyz5, x5, y5 = fou.plane_wave_propagation(Exyz4, x4, y4, dz=MCD.sequentialDistance[5], oversample=None,
show=False, save=False, outputCoordinates=None)
# From MCD-ADC-P3 to MCD-ADC-P4
print('Propagating from MCD-ADC-P3 to MCD-ADC-P4')
Exyz5 *= ELT_ADC.get_phase_delay(x5,y5,l, theta=None, paraxial=True, prism=3)
Exyz6, x6, y6 = fou.plane_wave_propagation(Exyz5, x5, y5, dz=MCD.sequentialDistance[6], oversample=None,
show=False, save=False, outputCoordinates=None)
# From MCD-ADC-P4 to MCD-CAM
print('Propagating from MCD-ADC-P4 to MCD-CAM')
Exyz6 *= ELT_ADC.get_phase_delay(x6,y6,l, theta=None, paraxial=True, prism=4)
Exyz7, x7, y7 = fou.plane_wave_propagation(Exyz6, x6, y6, dz=MCD.sequentialDistance[7], oversample=None,
show=False, save=False, outputCoordinates=None)
# From MCD-CAM to MCD-DET
print('Propagating from MCD-CAM to MCD-DET')
Exyz7 *= ELT_CAM.phase_delay(x7,y7,l, paraxial=True)
Exyz8, x8, y8 = fou.fresnel_propagation(Exyz7, x7, y7, dz=ELT_CAM.f, oversample=None,
show=False, save=False, outputCoordinates=None, apply_interpolation=False)
# # The result here shifts the PSF to the side, due to the phase tilts of the ADC
# # surfaces. We find the locations of the maximum and redo the the fresnel propagation
# # around these coordinates.
# print('Zooming in on the PSF location...')
# imax_1, imax_2 = find_maximum(np.abs(Exyz8))
# x8max = x8[imax_1, imax_2]
# y8max = y8[imax_1, imax_2]
# newCoordinates = np.asarray(coordinatesList[7]) + np.array([x8max, x8max, y8max, y8max])
# newCoordinates = newCoordinates.real.tolist()
# Exyz9, x9, y9 = fou.fresnel_propagation(Exyz7, x7, y7, dz=ELT_CAM.f, oversample=None,
# show=False, save=False, outputCoordinates=newCoordinates, apply_interpolation=True)
print('Propagation finished!')
###########################################################
######################## FITTING ##########################
###########################################################
print('Fitting the centroid...')
fitter = PSFFitter(x8.real,y8.real,np.abs(Exyz8)**2)
airyFitted, popt_airy, pcov_airy = fitter.fit_airy_PSF(initialGuess=(0.03**2, 0.0, 0.0, 0.0, 0.03))
gaussianFitted, popt_gauss, pcov_gauss = fitter.fit_gaussian_PSF(initialGuess=(0.04, 0, 0, 0.01, 0.01, 0, 0))
# moffatFitted, popt_moffat, pcov_moffat = fitter.fit_moffat_PSF(initialGuess=(0.04, 0, 0, 0.01, 0.01, 0, 0, 4))
custFitted, popt_cust, cust_var, cust_coupling = fitter.fit_custom_field(E0, Exyz8, return_coupling=True)
print('Strehl ratio: {:.4f}'.format(np.abs(cust_coupling)**2))
popt_airy_list.append(popt_airy)
popt_gauss_list.append(popt_gauss)
popt_cust_list.append(popt_cust)
pcov_airy_list.append([ np.sqrt(pcov_airy[i][i]) for i in range(len(pcov_airy)) ])
pcov_gauss_list.append([ np.sqrt(pcov_gauss[i][i]) for i in range(len(pcov_gauss)) ])
custom_var_list.append(np.sqrt(cust_var))
strehl_list.append(np.abs(cust_coupling)**2)
x0_cust, y0_cust = np.asarray(popt_cust).T
r_cust = np.sqrt(x0_cust**2 + y0_cust**2)
print (x0_cust, y0_cust)
# Save the variables
# save_variable('./data/wfe scaling/airy_5.pkl', [popt_airy_list, pcov_airy_list])
# save_variable('./data/wfe scaling/gauss_5.pkl', [popt_gauss_list, pcov_gauss_list])
# save_variable('./data/wfe scaling/cust_5.pkl', [popt_cust_list, custom_var_list, strehl_list])
# popt_airy_list, pcov_airy_list = load_variable('./data/wfe scaling/airy.pkl')
# popt_gauss_list, pcov_gauss_list = load_variable('./data/wfe scaling/gauss.pkl')
# popt_cust_list, custom_var_list, strehl_list = load_variable('./data/wfe scaling/cust.pkl')
amplitude_airy, x0_airy, y0_airy, R_airy, offset_airy = np.asarray(popt_airy_list).T
A_gauss, x0_gauss, y0_gauss, sigma_x_gauss, sigma_y_gauss, theta_gauss, offset_gauss = np.asarray(popt_gauss_list).T
# A_moffat, x0_moffat, y0_moffat, sigma_x_moffat, sigma_y_moffat, theta_moffat, offset_moffat, beta_moffat = np.asarray(popt_moffat_list).T
x0_cust, y0_cust = np.asarray(popt_cust_list).T # Check if this is correct!
# Errors in the optimum fit parameters for all scaled surface errors
amplitude_airy_err, x0_airy_err, y0_airy_err, R_airy_err, offset_airy_err = np.asarray(pcov_airy_list).T
A_gauss_err, x0_gauss_err, y0_gauss_err, sigma_x_gauss_err, sigma_y_gauss_err, theta_gauss_err, offset_gauss_err = np.asarray(pcov_gauss_list).T
# A_moffat_err, x0_moffat_err, y0_moffat_err, sigma_x_moffat_err, sigma_y_moffat_err, theta_moffat_err, offset_moffat_err, beta_moffat_err = np.asarray(pcov_moffat_list).T
x0_cust_err, y0_cust_err = np.asarray(custom_var_list).T
# Distance from original centroid at (0,0)
r_airy = np.sqrt(x0_airy**2+y0_airy**2)
r_err_airy = np.sqrt( (x0_airy**2*x0_airy_err**2+y0_airy**2*y0_airy_err**2) / (x0_airy**2+y0_airy**2))
r_gauss = np.sqrt(x0_gauss**2+y0_gauss**2)
r_err_gauss = np.sqrt( (x0_gauss**2*x0_gauss_err**2+y0_gauss**2*y0_gauss_err**2) / (x0_gauss**2+y0_gauss**2))
# r_moffat = np.sqrt(x0_moffat**2+y0_moffat**2)
# r_err_moffat = np.sqrt( (x0_moffat**2*x0_moffat_err**2+y0_moffat**2*y0_moffat_err**2) / (x0_moffat**2+y0_moffat**2))
r_cust = np.sqrt(x0_cust**2 + y0_cust**2)
r_err_cust = np.sqrt( (x0_cust**2*x0_cust_err**2+y0_cust**2*y0_cust_err**2) / (x0_cust**2+y0_cust**2))
###########################################################
######################## PLOTTING #########################
###########################################################
# Plot the PSF shape
# _xip = x8.real
# _yip = y8.real
# Exyz = Exyz8
# save = 'ELT-MCD-PSF_at_{}_WFE({}nm).pdf'.format((l*1000),(f*25))
# _fig = plt.figure(figsize=(16,6))
# _ax = _fig.add_subplot(121)
# _im = _ax.imshow(np.abs(Exyz.T)**2, cmap='hot',
# extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()], norm=LogNorm(vmin=1e-8))
# _ax.set_xlim([-10, 10])
# _ax.set_ylim([-10, 10])
# _ax.set_xlabel('x [mm]')
# _ax.set_ylabel('y [mm]')
# _fig.colorbar(_im, ax=_ax, label='Intensity')
# _ax2 = _fig.add_subplot(122)
# _im2 = _ax2.imshow(np.angle(np.angle(Exyz.T)), cmap='hot', extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()])
# _ax2.set_xlim([-10, 10])
# _ax2.set_ylim([-10, 10])
# _ax2.set_xlabel('x [mm]')
# _ax2.set_ylabel('y [mm]')
# _fig.colorbar(_im2, ax=_ax2, label='Phase')
# plt.tight_layout()
# plt.savefig(save, bbox_inches='tight', transparent=True)
# plt.show()
# Plot the centroid position
RMSlist = [25*f for f in fList]
fig = plt.figure()
ax1 = fig.add_subplot(121)
ax2 = fig.add_subplot(122)
ax1.errorbar(RMSlist, r_gauss, yerr=r_err_gauss, color='green', marker='o', linestyle='--', label='Gaussian', capsize=3, markeredgewidth=1)
ax1.errorbar(RMSlist, r_airy, yerr=r_err_airy, color='blue', marker='o', linestyle='--', label='Airy', capsize=3, markeredgewidth=1)
# ax1.errorbar(RMSlist, r_moffat, yerr=r_err_moffat, color='green', marker='o', linestyle='--', label='Moffat', capsize=3, markeredgewidth=1)
ax1.errorbar(RMSlist, r_cust, yerr=r_err_cust, color='orange', marker='o', linestyle='--', label='Custom', capsize=3, markeredgewidth=1)
# ax1.axhline(1.875e-4, color='red', linestyle='-', label='50 $\mu$as @ 4 mas/pix')
ax1.set_xscale('log')
ax1.set_yscale('log')
ax1.set_xlabel('RMS WFE (nm)')
ax1.set_ylabel('Distance')
ax1.legend()
ax1.grid(True, alpha=0.3, which='major')
ax1.grid(True, alpha=0.1, which='minor')
ax2.errorbar(x0_airy, y0_airy, xerr=x0_airy_err, yerr=y0_airy_err, color='k', marker='o', linestyle='--', label='Airy')
ax2.errorbar(x0_cust, y0_cust, xerr=x0_cust_err, yerr=y0_cust_err, color='r', marker='o', linestyle='--', label='Custom')
ax2.set_aspect('equal')
ax2.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
###########################################################
Reference in new issue View Git Blame Copy Permalink