def set_phase_map(self, z_coef, rho_zern=1., PV_rescale=True, PV_goal=50):
"""
Generates a phase map based on Zernike polynomials
This usually represents an NCPA map
Uses the fast methods implemented in zern_core.py
:param z_coef: coefficient of the Zernike series expansion
:param PV_goal: desired Peak-To-Valley [nm] of the phase map
"""
# Construct the base phase map
self.zern_model = zern.ZernikeNaive(mask=self.pupil_mask)
phase_map = self.zern_model(coef=z_coef,
rho=self.rho_m / rho_zern,
theta=self.theta_m,
normalize_noll=False,
mode='Jacobi',
print_option='Silent')
if PV_rescale:
# Compute the current PV and rescale the coefficients
current_pv = np.max(phase_map[np.nonzero(phase_map)]) - np.min(
phase_map[np.nonzero(phase_map)])
# Remember that at this point the phase map is a 1D array (a masked 2D)
phase_map = zern.rescale_phase_map(phase_map, peak=PV_goal / 2)
self.ncpa_coef = self.zern_model.coef * (PV_goal / current_pv
) # Save the coefficients
self.phase_map = zern.invert_mask(phase_map, self.pupil_mask)
else:
self.phase_map = zern.invert_mask(phase_map, self.pupil_mask)
self.ncpa_coef = self.zern_model.coef
def __init__(self, N_zern, initial_state):
### Zernike Wavefront
x = np.linspace(-1, 1, self.N_pix, endpoint=True)
xx, yy = np.meshgrid(x, x)
rho, theta = np.sqrt(xx**2 + yy**2), np.arctan2(xx, yy)
self.pupil = rho <= self.rho_aper
rho, theta = rho[self.pupil], theta[self.pupil]
zernike = zern.ZernikeNaive(mask=self.pupil)
_phase = zernike(coef=np.zeros(N_zern + 3),
rho=rho / self.rho_aper,
theta=theta,
normalize_noll=False,
mode='Jacobi',
print_option='Silent')
H_flat = zernike.model_matrix[:, 3:] # remove the piston and tilts
self.H_matrix = zern.invert_model_matrix(H_flat, self.pupil)
# Update the number of aberrations to match the dimensions of H
self.N = N_zern
self.N_zern = self.H_matrix.shape[-1]
self.PEAK = self.peak_PSF()
# Keep track of the STATE of the system
self.state = initial_state.copy()
def rand_zern(randker):
N = 48
N_zern = 10
rho_max = 0.9
eps_rho = 1.4
randgen = RandomState(randker)
extents = [-1, 1, -1, 1]
# Construct the coordinates
x = np.linspace(-rho_max, rho_max, N)
rho_spacing = x[1] - x[0]
xx, yy = np.meshgrid(x, x)
rho = np.sqrt(xx ** 2 + yy ** 2)
theta = np.arctan2(xx, yy)
aperture_mask = rho <= eps_rho * rho_max
rho, theta = rho[aperture_mask], theta[aperture_mask]
rho_max = np.max(rho)
extends = [-rho_max, rho_max, -rho_max, rho_max]
# Compute the Zernike series
coef = randgen.normal(size=N_zern)
z = zern.ZernikeNaive(mask=aperture_mask)
phase_map = z(coef=coef, rho=rho, theta=theta, normalize_noll=False, mode='Jacobi', print_option='Silent')
phase_map = zern.rescale_phase_map(phase_map, peak=1)
phase_2d = zern.invert_mask(phase_map, aperture_mask)
return phase_2d
def __init__(self, N_zern, initial_state, DM_stroke=0.05):
"""
Q-learning is based on DISCRETE actions. Therefore, if we want to correct for
Zernike aberrations which are continuous, we have to discretize the action space
:param N_zern: Number of Zernike aberrations we are expected to correct
:param DM_stroke: Deformable Mirror correction step in waves
"""
aberration_correction = [DM_stroke, -DM_stroke]
self.ACTION = N_zern * aberration_correction
### Initialize Zernike polynomials
x = np.linspace(-1, 1, self.N_pix, endpoint=True)
xx, yy = np.meshgrid(x, x)
rho, theta = np.sqrt(xx**2 + yy**2), np.arctan2(xx, yy)
self.pupil = rho <= self.rho_aper
rho, theta = rho[self.pupil], theta[self.pupil]
zernike = zern.ZernikeNaive(mask=self.pupil)
_phase = zernike(coef=np.zeros(N_zern + 3),
rho=rho / self.rho_aper,
theta=theta,
normalize_noll=False,
mode='Jacobi',
print_option='Silent')
H_flat = zernike.model_matrix[:, 3:] # remove the piston and tilts
self.H_matrix = zern.invert_model_matrix(H_flat, self.pupil)
# Update the number of aberrations to match the dimensions of H
self.N_zern = self.H_matrix.shape[-1]
self.PEAK = self.peak_PSF()
self.initial_state = initial_state
self.state = [initial_state]
self.rewards = [0]
def create_model_matrix(self, N_zern):
"""
Watch out because the matrices are in polar coordinates
:param N_zern:
:return:
"""
_coef = np.zeros(N_zern)
z = zern.ZernikeNaive(mask=np.ones((N, N)))
_phase = z(coef=_coef,
rho=self.rho,
theta=self.theta,
normalize_noll=False,
mode='Jacobi',
print_option=None)
self.model_matrix_flat = z.model_matrix
self.model_matrix = zern.invert_model_matrix(
z.model_matrix, np.ones((N, N)))
self.N_zern = self.model_matrix.shape[-1]
extents = [-1, 1, -1, 1]
# Construct the coordinates
x = np.linspace(-rho_max, rho_max, N)
rho_spacing = x[1] - x[0]
xx, yy = np.meshgrid(x, x)
rho = np.sqrt(xx**2 + yy**2)
theta = np.arctan2(xx, yy)
aperture_mask = rho <= eps_rho * rho_max
rho, theta = rho[aperture_mask], theta[aperture_mask]
rho_max = np.max(rho)
extends = [-rho_max, rho_max, -rho_max, rho_max]
# Compute the Zernike series
coef = randgen.normal(size=N_zern)
z = zern.ZernikeNaive(mask=aperture_mask)
phase_map = z(coef=coef,
rho=rho,
theta=theta,
normalize_noll=False,
mode='Jacobi',
print_option=None)
# phase_map = zern.rescale_phase_map(phase_map, peak=1)
phase_2d = zern.invert_mask(phase_map, aperture_mask)
# Introduce some noise in the map
noised_phase_map = phase_map + 0.5 * np.random.normal(size=phase_map.shape[0])
noised_2d = zern.invert_mask(noised_phase_map, aperture_mask)
plt.figure()
if __name__ == "__main__":
x = np.linspace(-Luv / 2, Luv / 2, N_pix, endpoint=True)
xx, yy = np.meshgrid(x, x)
r = np.sqrt(xx**2 + yy**2)
theta = np.arctan2(xx, yy)
pupil_mask = xx**2 + yy**2 <= rho**2
r, theta = r[pupil_mask], theta[pupil_mask]
# Initialize the ZernikeNaive to get access to the Model Matrix
np.random.seed(123)
coef = np.random.normal(size=N_zern)
z = zern.ZernikeNaive(mask=pupil_mask)
_phase = z(coef=coef,
rho=r / rho_zern,
theta=theta,
normalize_noll=False,
mode='Jacobi',
print_option='Silent')
# ==================================================================================================================
""" Model matrix for Zernike polynomials """
# H contains the Zernike Polynomials
H = reshape_model_matrix(z.model_matrix, pupil_mask)
N_zern = H.shape[-1]
print('Using %d Zernikes' % N_zern)
def evaluate_wavefront_performance(N_zern,
test_coef,
guessed_coef,
zern_list,
twisted=False,
show_predic=False):
"""
Evaluates the performance of the ML method regarding the final
RMS wavefront error. Compares the initial RMS NCPA and the residual
after correction
"""
# Transform the ordering to match the Zernike matrix
new_test_coef = transform_zemax_to_noll(test_coef, twisted=False)
new_guessed_coef = transform_zemax_to_noll(guessed_coef, twisted)
x = np.linspace(-1, 1, 512, endpoint=True)
xx, yy = np.meshgrid(x, x)
rho, theta = np.sqrt(xx**2 + yy**2), np.arctan2(xx, yy)
pupil = rho <= 1.0
rho, theta = rho[pupil], theta[pupil]
zernike = zern.ZernikeNaive(mask=pupil)
_phase = zernike(coef=np.zeros(new_test_coef.shape[1] + 3),
rho=rho,
theta=theta,
normalize_noll=False,
mode='Jacobi',
print_option='Silent')
H_flat = zernike.model_matrix[:, 3:] # remove the piston and tilts
H_matrix = zern.invert_model_matrix(H_flat, pupil)
# print(H_flat.shape)
# Elliptical mask
ellip_mask = (xx / 0.5)**2 + (yy / 1.)**2 <= 1.0
H_flat = H_matrix[ellip_mask]
# print(H_flat.shape)
N = test_coef.shape[0]
initial_rms = np.zeros(N)
residual_rms = np.zeros(N)
for k in range(N):
phase = np.dot(H_flat, new_test_coef[k])
residual_phase = phase - np.dot(H_flat, new_guessed_coef[k])
before, after = np.std(phase), np.std(residual_phase)
initial_rms[k] = before
residual_rms[k] = after
average_initial_rms = np.mean(initial_rms)
average_residual_rms = np.mean(residual_rms)
improvement = (average_initial_rms -
average_residual_rms) / average_initial_rms * 100
print('\nWAVEFRONT PERFORMANCE DATA')
print('\nNumber of samples in TEST dataset: %d' % N)
print('Average INITIAL RMS: %.3f waves (%.1f nm @1.5um)' %
(average_initial_rms, average_initial_rms * wave_nom))
print('Average RESIDUAL RMS: %.3f waves (%.1f nm @1.5um)' %
(average_residual_rms, average_residual_rms * wave_nom))
print('Improvement: %.2f percent' % improvement)
if show_predic == True:
plt.figure()
plt.scatter(range(N),
initial_rms * wave_nom,
c='red',
s=6,
label='Before')
plt.scatter(range(N),
residual_rms * wave_nom,
c='blue',
s=6,
label='After')
plt.xlabel('Test PSF')
plt.xlim([0, N])
plt.ylim(bottom=0)
plt.ylabel('RMS wavefront [nm]')
# plt.title(r'$\lambda=1.5$ $\mu$m (defocus: 0.20 waves)')
plt.legend(title='Calibration stage')
N_ok = (np.argwhere(residual_rms * wave_nom < 100)).shape[0]
plt.figure()
plt.scatter(initial_rms * wave_nom,
residual_rms * wave_nom,
c='blue',
s=8)
plt.axhline(y=100, linestyle='--')
plt.xlabel('Initial RMS [nm]')
plt.ylabel('Residual RMS [nm]')
plt.title('%d / %d cases with RMS < 100 nm' % (N_ok, N))
plt.ylim(bottom=0)
plt.figure()
n_bins = 20
for k in range(N_zern):
guess = guessed_coef[:, k]
coef = test_coef[:, k]
residual = coef - guess
mu, s2 = np.mean(residual), (np.std(residual))
label = zern_list[k] + r' ($\mu$=%.3f, $\sigma$=%.2f)' % (mu, s2)
plt.hist(residual, histtype='step', label=label)
plt.legend(title=r'Residual aberrations [waves]', loc=2)
plt.xlabel(r'Residual [waves]')
plt.xlim([-0.075, 0.075])
for k in range(N_zern):
guess = guessed_coef[:, k]
coef = test_coef[:, k]
colors = wave_nom * residual_rms
colors -= colors.min()
colors /= colors.max()
colors = cm.rainbow(colors)
plt.figure()
ss = plt.scatter(coef, guess, c=colors, s=20)
x = np.linspace(-0.10, 0.10, 10)
# plt.colorbar(ss)
plt.plot(x, x, color='black', linestyle='--')
title = zern_list[k]
plt.title(title)
plt.xlabel('True Value [waves]')
plt.ylabel('Predicted Value [waves]')
plt.xlim([-0.10, 0.10])
plt.ylim([-0.10, 0.10])
return initial_rms, residual_rms
""" import numpy as np import matplotlib.pyplot as plt import zern_core as zern # Parameters N = 1024 N_zern = 100 rho_max = 1.0 mask = [] # Construct the coordinates rho = np.linspace(0., rho_max, N) Z_naive = zern.ZernikeNaive(mask) # Zernike orders n_list = [40, 44, 48] m = 0 n_max = 101 # for n in n_list: # # Compute the Radial Zernike R_nm # R_naive = Z_naive.R_nm(n, m, rho) # R_jacobi = Z_naive.R_nm_Jacobi(n, m, rho) # # plt.figure() # plt.plot(rho, R_naive, label='Standard') # plt.plot(rho, R_jacobi, label='Jacobi') # plt.legend()
randgen = RandomState(1234)
coef = randgen.normal(size=N_zern)
print('First 10 Zernike coefficients')
print(coef[:10])
# --------------------------------------------------------------------
""" Plot the Wavefront in Polar coordinates """
rho_1 = np.linspace(0.0, 1.0, N, endpoint=True)
theta_1 = np.linspace(0.0, 2 * np.pi, N)
rr, tt = np.meshgrid(rho_1, theta_1)
rho, theta = rr.flatten(), tt.flatten()
z = zern.ZernikeNaive(mask=np.ones((N, N)))
phase = z(coef=coef, rho=rho, theta=theta, normalize_noll=False, mode='Jacobi', print_option=None)
phase_map = phase.reshape((N, N))
fig = plt.figure()
ax = fig.add_subplot(111)
ax.imshow(phase_map, origin='lower', cmap='viridis', extent=[0, 1, 0, 2 * np.pi])
ax.set_aspect('auto')
plt.xlabel(r'Radius $\rho$')
plt.ylabel(r'Angle $\theta$')
# plt.title('Wavefront map in Polar Coordinates')
# --------------------------------------------------------------------
""" Plot the Wavefront in Cartesian coordinates """
xx = rr * np.cos(tt)
path_prelim = os.path.abspath(
'D:/Thesis/LAM/POP/Slicer/0 Preliminary Tasks')
N_slices = len(list_slices)
# ===================
x = np.linspace(-1, 1, 1024, endpoint=True)
xx, yy = np.meshgrid(x, x)
rho_circ, theta_circ = np.sqrt((xx)**2 + yy**2), np.arctan2(xx, yy)
rho_elli, theta_elli = np.sqrt((xx)**2 + (2 * yy)**2), np.arctan2(xx, yy)
pupil_circ = rho_circ <= 1.0
pupil_elli = rho_elli <= 1.0
### Clipped Defocus
rho_circ, theta_circ = rho_circ[pupil_circ], theta_circ[pupil_circ]
zernike = zern.ZernikeNaive(mask=pupil_circ)
_phase = zernike(coef=np.zeros(50),
rho=rho_circ,
theta=theta_circ,
normalize_noll=False,
mode='Jacobi',
print_option='Silent')
H_flat = zernike.model_matrix # remove the piston and tilts
H_matrix = zern.invert_model_matrix(H_flat, pupil_circ)
defocus_circ = H_matrix[:, :, 4].copy()
### Elliptic
rho_elli, theta_elli = rho_elli[pupil_elli], theta_elli[pupil_elli]
zernike = zern.ZernikeNaive(mask=pupil_elli)
_phase = zernike(coef=np.zeros(25),
rho=rho_elli,
spaxel_scale = 0.5 # mas
RHO_APER = rho_spaxel_scale(spaxel_scale, wavelength=wave)
check_spaxel_scale(RHO_APER, wave)
FWHM = compute_FWHM(wave) # mas
crop_pix = 5 * FWHM / spaxel_scale # How many pixels to cover enough FWHMs
crop_pix = int(np.ceil(crop_pix))
""" Impact of Defocus on a Round PSF - Evolution of Rings """
# How does the defocus intensity affect the Peak and Rings of the PSF?
x = np.linspace(-1, 1, N_PIX, endpoint=True)
xx, yy = np.meshgrid(x, x)
rho, theta = np.sqrt(xx**2 + (2 * yy)**2), np.arctan2(xx, yy)
pupil = rho <= RHO_APER
rho, theta = rho[pupil], theta[pupil]
zernike = zern.ZernikeNaive(mask=pupil)
_phase = zernike(coef=np.zeros(10),
rho=rho / RHO_APER,
theta=theta,
normalize_noll=False,
mode='Jacobi',
print_option='Silent')
H_flat = zernike.model_matrix[:, 3:] # remove the piston and tilts
H_matrix = zern.invert_model_matrix(H_flat, pupil)
# Rescale by 1/2 to have a Peak-To-Valley of 1 wave
defocus = 0.5 * H_matrix[:, :, 1].copy()
# Compute nominal PSF -> Pupil = Aperture * exp(2 PI wavefront)
pupil_nom = pupil * np.exp(2 * np.pi * 1j * np.zeros_like(pupil))
image_nom = (np.abs(fftshift(fft2(pupil_nom))))**2
print('N_pixels: %d' % N_pix)
print('D_ELT/L = %.2f' % eps)
print('Wavelenght: %.1f [microns]' % wave)
print('T exposure: %.2f' % t_exp)
print('Max Zernikes: %d' % max_N_zern)
print('====================================================\n')
pupil_mask = mt.elt_mask(eps, xx, yy)
# WATCH OUT because this rho_masked is renormalized to 1
# so that the zern_model.model_matrix follows the convention of giving
# Z = 1 at the borders (i.e the typical Zernike map)
rho_masked, theta_masked = rho[pupil_mask] / eps, theta[pupil_mask]
# Initialize the Zernike Model
zern_model = zern.ZernikeNaive(mask=pupil_mask)
# Use a Random vector to run it once and thus
# initiliaze the Model Matrix H
_coef = np.random.normal(size=max_N_zern)
zern_model(_coef,
rho_masked,
theta_masked,
normalize_noll=False,
mode='Jacobi',
print_option='Silent')
H_matrix = zern_model.model_matrix
""" ++++++++++++++++++++++++++++++++++++++++++++++++++++++ """
""" MACHINE LEARNING """
""" ++++++++++++++++++++++++++++++++++++++++++++++++++++++ """
smart_psf = fz.SmartPSF(zern_model=zern_model,
# Compare the IDEAL and the TRUE PSFs
plt.figure()
plt.subplot(121)
plt.imshow(nom_psf, extent=extends, cmap='jet')
plt.xlabel('X [mm]')
plt.ylabel('Y [mm]')
plt.title('Nominal PSF (Slicer)')
plt.subplot(122)
plt.imshow(pop.crop_array(image, N_pix), extent=extends, cmap='jet')
plt.xlabel('X [mm]')
plt.title('Ideal PSF')
# Show an example of NCPA map
zern_model = zern.ZernikeNaive(mask=aper_mask)
zern_coef = np.zeros(10)
zern_coef[4] = 0.25
zz = zern_model(coef=zern_coef,
rho=rho,
theta=theta,
mode='Standard',
print_option=None)
phase = zern.invert_mask(zz, aper_mask)
H = zern_model.model_matrix
plt.figure()
plt.imshow(zern.invert_mask(H[:, 4], aper_mask), extent=[-1, 1, -1, 1])
plt.xlim([-1.1 * eps, 1.1 * eps])
plt.ylim([-1.1 * eps, 1.1 * eps])