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 __init__(self,
xx,
yy,
pupil_mask,
rho_aper,
rho_zern,
images,
phase_diversity,
zern_model,
wave,
t_exp=1.0):
self.t_exp = t_exp
self.pupil_mask = pupil_mask
self.N_images = len(images)
self.N_pix = xx.shape[0]
self.image_nominal = images[0]
self.image_plus_defocus = images[1]
if self.N_images == 3:
self.image_minus_defocus = images[2]
self.wave = wave
self.phase_diversity = phase_diversity
self.zern_model = zern_model
self.rho_aper, self.rho_zern = rho_aper, rho_zern
# Take of the coordinates and masked arrays
rho, theta = np.sqrt(xx**2 + yy**2), np.arctan2(xx, yy)
self.rho_m, self.theta_m = rho[self.pupil_mask], theta[self.pupil_mask]
# Reshape the Zern Model Matix H
H = zern_model.model_matrix
self.model_matrix = zern.invert_model_matrix(H, pupil_mask)
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 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]
def __init__(self, images, eps, aper_mask, zern_model, phase_diversity):
self.images = images
self.Npix = images[0].shape[0]
self.aper_mask = aper_mask
self.n_crop = self.aper_mask.shape[0] // self.Npix
self.zern_model = zern_model
self.defocus = phase_diversity
# Reshape the Zern Model Matix H
H = zern_model.model_matrix
self.model_matrix = zern.invert_model_matrix(H, self.aper_mask)
x = np.linspace(-1., 1., self.aper_mask.shape[0], endpoint=True)
xx, yy = np.meshgrid(x, x)
rho = np.sqrt(xx**2 + yy**2)
theta = np.arctan2(yy, xx)
rho, theta = rho[self.aper_mask], theta[self.aper_mask]
rho /= eps
self.rho, self.theta = rho, theta
self.compute_normalization()
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
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,
theta=theta_elli,
normalize_noll=False,
mode='Jacobi',
print_option='Silent')
H_flat_elli = zernike.model_matrix # remove the piston and tilts
H_matrix_elli = zern.invert_model_matrix(H_flat_elli, pupil_elli)
piston_elli = H_matrix_elli[:, :, 0].copy()
# or whether you want it row_by_row == True
# This is because low order and high order aberrations behave differently wrt EE
row_by_row = False
for level in np.arange(4, 11, 2):
print("\nZernike Level: ", level)
# Calculate how many Zernikes we need to have up to a certain Radial level
exp = level - 1 # Exponent of the Zernike radial polynomial
N_zern = triang[level] # How many Zernikes in total to ask zern to use
_coef = np.zeros(N_zern)
z = zern.ZernikeNaive(mask=pupil)
_phase = z(coef=_coef, rho=rho/RHO_APER, theta=theta, normalize_noll=False, mode='Jacobi',
print_option=None)
model_matrix_flat = z.model_matrix
model_matrix = zern.invert_model_matrix(z.model_matrix, pupil)
model_matrix = model_matrix[:, :, 3:] # remove piston and tilts
model_matrix_flat = model_matrix_flat[:, 3:]
if row_by_row: # Select only a specific row (radial order) of Zernike
# Number of polynomials in that last Zernike row
poly_in_row = triang[level] - triang[level - 1]
model_matrix = model_matrix[:, :, -poly_in_row:] # select only the high orders
model_matrix_flat = model_matrix_flat[:, -poly_in_row:]
N_zern = model_matrix.shape[-1] # Update the N_zern after removing Piston and Tilts
PSF_zern = PointSpreadFunction([model_matrix, pupil, model_matrix_flat])
plt.figure()
# 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
PEAK = np.max(image_nom)
image_nom /= PEAK
focus = [0, 0.25, 0.5, 1.0]
img = []
mas = spaxel_scale * np.linspace(-crop_pix // 2, crop_pix // 2, crop_pix)
plt.figure()
for eps in focus:
