mirror of
https://github.com/yosata/fourierOptics.git
synced 2026-10-29 20:51:10 +00:00
408 lines
16 KiB
Python
408 lines
16 KiB
Python
#!/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 sys
|
|
sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Documents/Published/FourierOptics GitHub release/')
|
|
from fourierOptics import *
|
|
|
|
|
|
'''
|
|
@ Author: Joost van den Born
|
|
@ Desc.: Importing a grid sag file and scaling it appropriately
|
|
2022-08-02: Old script. Currently broken.
|
|
'''
|
|
|
|
class surface:
|
|
def __init__(self,data):
|
|
self.data = data
|
|
self.resolution = np.shape(self.data)
|
|
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))
|
|
|
|
def scale_rms(self,factor):
|
|
# Scale the surface data such that the RMS gets multiplied by factor
|
|
_factor = factor ** 2
|
|
self.calculate_2D_PSD()
|
|
self.scaledPSD2D = self.PSD2D * _factor
|
|
# Fourier transform of the surface profile
|
|
_in = pyfftw.empty_aligned((n,m), dtype='complex128')
|
|
_in[:] = self.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()
|
|
_phase = np.angle(_out)
|
|
|
|
_in = pyfftw.empty_aligned((n,m), dtype='complex128')
|
|
_in[:] = np.exp(1j*_phase) * np.sqrt(np.fft.fftshift(self.scaledPSD2D))
|
|
_out = pyfftw.empty_aligned((n,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))
|
|
|
|
def calculate_2D_PSD(self):
|
|
''' Returns the 2D PSD map '''
|
|
_in = pyfftw.empty_aligned((n,m), dtype='complex128')
|
|
_in[:] = self.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()
|
|
self.D1 = np.fft.fftshift(_out)
|
|
self.PSD2D = np.abs(self.D1)**2
|
|
return self.PSD2D
|
|
|
|
def calculate_1D_PSD(self, filteredDataFlag=False):
|
|
''' Returns the azimuthally averaged PSD map '''
|
|
try:
|
|
if not filteredDataFlag:
|
|
self.PSD1D = self.azimuthal_average()
|
|
else:
|
|
self.filteredPSD1D = self.azimuthal_average(filteredDataFlag=filteredDataFlag)
|
|
return self.filteredPSD1D
|
|
except AttributeError:
|
|
self.calculate_2D_PSD()
|
|
self.PSD1D = self.azimuthal_average()
|
|
|
|
return self.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
|
|
_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
|
|
|
|
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]
|
|
_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
|
|
|
|
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]
|
|
_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
|
|
|
|
def azimuthal_average(self, center=None, filteredDataFlag=False):
|
|
"""
|
|
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:
|
|
_i_sorted = self.PSD2D.flat[_ind]
|
|
else:
|
|
_i_sorted = self.filteredPSD2D.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
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
plt.rc('font', family='serif')
|
|
plt.rc('axes', labelsize='large')
|
|
|
|
# Pick a wavelength
|
|
l = 1.35e-3
|
|
|
|
# Initiate Fourier Optics
|
|
fou = FourierOptics(l=l)
|
|
|
|
# Pick sampling parameters
|
|
dpix = 160
|
|
n, m = 1001, 1001
|
|
oC = [-.5, .5, -.5, .5]
|
|
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.top_hat(38854/2, 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))
|
|
# oC = [x.min(), x.max(), y.min(), y.max()]
|
|
|
|
# Add PSD wavefront error to the lens
|
|
wfe = WavefrontError()
|
|
h_data = wfe.generate_height_map(5e-5, 2, 1, n, m) # 50 nm RMS for standard
|
|
|
|
PSDSurface = surface(h_data)
|
|
# print(PSDSurface.rms)
|
|
# PSDSurface.set_bandpass_filter(10, 20)
|
|
# PSDSurface.calculate_1D_PSD(filteredDataFlag=False)
|
|
# PSDSurface.calculate_1D_PSD(filteredDataFlag=True)
|
|
|
|
# fig = plt.figure()
|
|
# ax1 = fig.add_subplot(121)
|
|
# ax2 = fig.add_subplot(122)
|
|
# im1 = ax1.imshow(PSDSurface.filteredData, origin='lower', cmap='jet')
|
|
# ax2.loglog(PSDSurface.PSD1D)
|
|
# ax2.loglog(PSDSurface.filteredPSD1D)
|
|
# ax2.grid(alpha=0.3, which='major')
|
|
# ax2.grid(alpha=0.1, which='minor')
|
|
# fig.colorbar(im1, ax=ax1)
|
|
# plt.tight_layout()
|
|
# plt.show()
|
|
|
|
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 = []
|
|
Exyzlist = []
|
|
# Scaling the surface profile by a factor f
|
|
|
|
|
|
|
|
|
|
'''
|
|
|
|
|
|
M1 = ParabolicMirror(dx=0, dy=0, name='ELT-M1', CA=38854/2, f=679945)
|
|
t_l = M1.phase_delay(x,y,l)
|
|
|
|
# The input distribution gets a curved wavefront after the lens
|
|
U_in = A * t_l
|
|
|
|
Exyz_og,xp,yp = fou.fresnel_propagation(U_in, x, y, dz=M1.f, show=False, outputCoordinates=oC)
|
|
|
|
# Preparing some figures
|
|
fig1 = plt.figure()
|
|
ax = fig1.add_subplot(111)
|
|
im = ax.imshow(np.abs(Exyz_og)**2, cmap=plt.cm.jet, origin='bottom',
|
|
extent=oC)
|
|
fig1.colorbar(im)
|
|
|
|
fig2 = plt.figure()
|
|
ax21 = fig2.add_subplot(121)
|
|
ax21.semilogy((np.abs(Exyz_og)**2)[500], 'k-', label='Original')
|
|
ax21.grid(True, alpha=0.3)
|
|
|
|
|
|
ax22 = fig2.add_subplot(122)
|
|
ax22.grid(True, alpha=0.3)
|
|
|
|
|
|
fig3 = plt.figure()
|
|
ax31 = fig3.add_subplot(131)
|
|
ax32 = fig3.add_subplot(132)
|
|
ax33 = fig3.add_subplot(133)
|
|
|
|
# Start of fitting routine
|
|
psfFit = PSFFitter(xp.real,yp.real,np.abs(Exyz_og)**2)
|
|
|
|
gaussianFit = moffatFit = airyFit = True
|
|
|
|
# Fit a Gaussian profile
|
|
if gaussianFit:
|
|
gaussianFitted, popt_gauss, pcov_gauss = psfFit.fit_gaussian_PSF(initialGuess=(0.03**2, 0, 0, 0.02, 0.02, 0, 0))
|
|
A_gauss, x0_gauss, y0_gauss, sigma_x_gauss, sigma_y_gauss, theta_gauss, offset_gauss = popt_gauss
|
|
ax.contour(xp.real, yp.real, gaussianFitted, 8, colors='w', linestyles='solid')
|
|
ax.plot(x0_gauss,y0_gauss, 'wx')
|
|
ax21.semilogy(gaussianFitted[500], 'r-', label='Gaussian fit')
|
|
ax22.semilogy((np.abs(Exyz_og)**2)[500] - gaussianFitted[500], 'r-', label='Gaussian differential')
|
|
im31 = ax31.imshow(np.abs(Exyz_og)**2-gaussianFitted, cmap=plt.cm.jet, origin='bottom', extent=oC)
|
|
fig3.colorbar(im31, ax=ax31)
|
|
|
|
if moffatFit:
|
|
moffatFitted, popt_moffat, pcov_moffat = psfFit.fit_moffat_PSF(initialGuess=(0.03**2, 0, 0, 0.02, 0.02, 0, 0, 4))
|
|
A_moffat, x0_moffat, y0_moffat, sigma_x_moffat, sigma_y_moffat, theta_moffat, offset_moffat, beta_moffat = popt_moffat
|
|
ax.contour(xp.real, yp.real, moffatFitted, 8, colors='y', linestyles='dashed')
|
|
ax.plot(x0_moffat,y0_moffat, 'yx')
|
|
ax21.semilogy(moffatFitted[500], 'g-', label='Moffat fit')
|
|
ax22.semilogy((np.abs(Exyz_og)**2)[500] - moffatFitted[500], 'g-', label='Moffat differential')
|
|
im32 = ax32.imshow(np.abs(Exyz_og)**2-moffatFitted, cmap=plt.cm.jet, origin='bottom', extent=oC)
|
|
fig3.colorbar(im32, ax=ax32)
|
|
|
|
if airyFit:
|
|
airyFitted, popt_airy, pcov_airy = psfFit.fit_airy_PSF(initialGuess=(0.03**2, 0.0, 0.0, 0.0, 0.03))
|
|
amplitude_airy, x0_airy, y0_airy, R_airy, offset_airy = popt_airy
|
|
ax.contour(xp.real, yp.real, airyFitted, 8, colors='c', linestyles='dashdot')
|
|
ax.plot(x0_airy,y0_airy, 'cx')
|
|
ax21.semilogy(airyFitted[500], 'b-', label='Airy fit')
|
|
ax22.semilogy((np.abs(Exyz_og)**2)[500] - airyFitted[500], 'b-', label='Airy differential')
|
|
im33 = ax33.imshow(np.abs(Exyz_og)**2-airyFitted, cmap=plt.cm.jet, origin='bottom', extent=oC)
|
|
fig3.colorbar(im33, ax=ax33)
|
|
|
|
ax21.legend()
|
|
ax22.legend()
|
|
plt.tight_layout()
|
|
plt.show()
|
|
|
|
|
|
|
|
|
|
'''
|
|
|
|
|
|
# Create a lens with a focal length ~100 mm, with a clear aperture radius of 4 mm and thickness of 1 mm.
|
|
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
|
|
Exyz_og,xp,yp = fou.fresnel_propagation(U_in, x, y, dz=M1.f, show=False, outputCoordinates=oC)
|
|
|
|
eef2, r2 = fou.encircled_energy_function(np.abs(Exyz_og)**2, xp, yp)
|
|
plt.plot(r2, eef2); plt.grid();plt.show()
|
|
set_trace()
|
|
RMSlist = []
|
|
flist = [0.01, 0.1, 0.2, 0.5, 0.75, 1, 1.5, 2, 3, 4, 5, 6]
|
|
for f in flist:
|
|
PSDSurface.scale_rms(f)
|
|
scal = PSDSurface.scaledData
|
|
print(f, '{:.1f} nm'.format(PSDSurface.scaledrms*1e6))
|
|
RMSlist.append(PSDSurface.scaledrms*1e6)
|
|
M1.add_wavefront_error(scal)
|
|
t_l = M1.phase_delay(x,y,l)
|
|
|
|
# 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.
|
|
Exyz_psd,xp,yp = fou.fresnel_propagation(U_in, x, y, dz=M1.f, show=False, outputCoordinates=oC)
|
|
|
|
# Set up fitting object
|
|
psfFit_psd = PSFFitter(xp.real,yp.real,np.abs(Exyz_psd)**2)
|
|
|
|
# Fit airy pattern, Gaussian and Moffat functions
|
|
airyFitted, popt_airy, pcov_airy = psfFit_psd.fit_airy_PSF(initialGuess=(0.03**2, 0.0, 0.0, 0.0, 0.03))
|
|
gaussianFitted, popt_gauss, pcov_gauss = psfFit_psd.fit_gaussian_PSF(initialGuess=(0.04, 0, 0, 0.01, 0.01, 0, 0))
|
|
# moffatFitted, popt_moffat, pcov_moffat = psfFit_psd.fit_moffat_PSF(initialGuess=(0.04, 0, 0, 0.01, 0.01, 0, 0, 4))
|
|
custFitted, popt_cust, cust_var, cust_coupling = psfFit_psd.fit_custom_field(Exyz_og, Exyz_psd, return_coupling=True)
|
|
|
|
|
|
popt_airy_list.append(popt_airy)
|
|
popt_gauss_list.append(popt_gauss)
|
|
# popt_moffat_list.append(popt_moffat)
|
|
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)) ])
|
|
# pcov_moffat_list.append([ np.sqrt(pcov_moffat[i][i]) for i in range(len(pcov_moffat)) ])
|
|
custom_var_list.append(np.sqrt(cust_var))
|
|
strehl_list.append(np.abs(cust_coupling)**2)
|
|
Exyzlist.append(Exyz_psd)
|
|
|
|
|
|
|
|
# Optimum fit parameters for all scaled surface errors
|
|
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))
|
|
|
|
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(block=False)
|
|
|
|
fig2 = plt.figure()
|
|
for i in range(len(flist)):
|
|
axi = fig2.add_subplot(4,3,i+1)
|
|
imi = axi.imshow(np.abs(Exyzlist[i])**2, extent=oC, cmap='jet')
|
|
# fig2.colorbar(imi, ax=axi)
|
|
axi.title.set_text('RMS: {:.1f} nm (S={:.2f})'.format(RMSlist[i], strehl_list[i]))
|
|
axi.set_xlim([-0.2,0.2])
|
|
axi.set_ylim([-0.2,0.2])
|
|
plt.tight_layout()
|
|
plt.show()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|