def rescale_coefficients_for_strehl(PSF_model,
coef_strength,
rescale,
min_strehl=0.25,
a=0.5,
b=0.75):
# Calculate dummy datasets
a_PSF, _coef, _PSF, __coef = calibration.generate_dataset(
PSF_model, 99, 1, a * coef_strength, a * rescale)
b_PSF, _coef, _PSF, __coef = calibration.generate_dataset(
PSF_model, 99, 1, b * coef_strength, b * rescale)
# Calculate peaks of training sets
peaks_a = np.max(a_PSF[:, :, :, 0], axis=(1, 2))
peaks_b = np.max(b_PSF[:, :, :, 0], axis=(1, 2))
# Calculate the average minimum Strehl
min_a = np.mean(np.sort(peaks_a)[:5])
min_b = np.mean(np.sort(peaks_b)[:5])
# Calculate rescaling of coefficients necessary to have a min_strehl
ab = b - a
dstrehl = min_b - min_a
ds = min_strehl - min_a
da = ab * ds / dstrehl
return a + da
div_phase = PSF_actuators.diversity_phase
rms_foc = np.std(div_phase[PSF_actuators.pupil_mask])
print("\nRMS Diversity: %.2f" % rms_foc)
# rms_div.append(rms_foc)
rms_div_noisy[k_noise, i_beta] = rms_foc
# plt.figure()
# plt.imshow(PSF_actuators.diversity_phase, cmap='RdBu')
# plt.colorbar()
# plt.title(r'Diversity Map | Defocus [rad]')
# Generate a training set
train_PSF, train_coef, test_PSF, test_coef = calibration.generate_dataset(
PSF_actuators,
N_train,
N_test,
coef_strength=coef_strength,
rescale=rescale)
# Add some Readout Noise to the PSF images to spice things up
SNR = noiseSNR[k_noise]
noise_model = noise.NoiseEffects()
train_PSF_readout = noise_model.add_readout_noise(train_PSF,
RMS_READ=1. /
SNR)
test_PSF_readout = noise_model.add_readout_noise(test_PSF,
RMS_READ=1. / SNR)
#
# # Show the Defocus phase, and some PSF examples
# ax1 = axes[i_beta][0]
# img1 = ax1.imshow(div_phase, cmap='RdBu', extent=[-1, 1, -1, 1])
# (1) We begin by creating a Zernike PSF model with Defocus as diversity
zernike_matrix, pupil_mask_zernike, flat_zernike = psf.zernike_matrix(N_levels=N_levels, rho_aper=RHO_APER,
rho_obsc=RHO_OBSC,
N_PIX=N_PIX, radial_oversize=1.0,
anamorphic_ratio=1.0)
N_zern = zernike_matrix.shape[-1]
zernike_matrices = [zernike_matrix, pupil_mask_zernike, flat_zernike]
PSF_zernike = psf.PointSpreadFunction(matrices=zernike_matrices, N_pix=N_PIX,
crop_pix=pix, diversity_coef=np.zeros(zernike_matrix.shape[-1]))
defocus_zernike = np.zeros(zernike_matrix.shape[-1])
defocus_zernike[1] = diversity / (2 * np.pi)
PSF_zernike.define_diversity(defocus_zernike)
# C
train_PSF, train_coef, test_PSF, test_coef = calibration.generate_dataset(PSF_zernike, N_train, N_test,
coef_strength, rescale)
# Train the Calibration Model on images with the nominal defocus
epochs = 10
calib_zern = calibration.Calibration(PSF_model=PSF_zernike)
calib_zern.create_cnn_model(layer_filters, kernel_size, name='NOM_ZERN', activation='relu')
losses = calib_zern.train_calibration_model(train_PSF, train_coef, test_PSF, test_coef,
N_loops=1, epochs_loop=epochs, verbose=1, batch_size_keras=32, plot_val_loss=False,
readout_noise=False, RMS_readout=[1. / SNR], readout_copies=0)
# Now we test the performance on an anamorphic error
ratio = 1.10
zernike_matrix_anam, pupil_mask_zernike_anam, flat_zernike_anam = psf.zernike_matrix(N_levels=N_levels, rho_aper=RHO_APER,
rho_obsc=RHO_OBSC, N_PIX=N_PIX,
radial_oversize=1.0,
anamorphic_ratio=ratio)
]
extent_anam = [-pix // 2 * SPAX, pix // 2 * SPAX, -pix * SPAX, pix * SPAX]
fig, (ax1, ax2) = plt.subplots(1, 2)
img1 = ax1.imshow(p_nom, cmap=cmap, extent=extent)
img2 = ax2.imshow(p_anam, cmap=cmap, extent=extent_anam)
ax2.set_ylim([-pix // 2 * SPAX, pix // 2 * SPAX])
plt.show()
# ================================================================================================================ #
# FIELD OF VIEW investigation
# ================================================================================================================ #
# Does the number of pixels used to crop the images matter?
train_PSF, train_coef, test_PSF, test_coef = calibration.generate_dataset(
PSF_zernike, N_train, N_test, coef_strength, rescale)
def crop_images(PSF_images, crop_pix):
"""
Get some PSF datacubes and crop them
:param PSF_images:
:param crop_pix:
:return:
"""
N_PSF = PSF_images.shape[0]
N_pix = PSF_images.shape[1]
min_pix = N_pix // 2 - crop_pix // 2
max_pix = N_pix // 2 + crop_pix // 2
new_images = np.zeros((N_PSF, crop_pix, crop_pix, 2))
N_alphas = 20
alphas = np.logspace(0, 2, N_alphas)
strehls = np.zeros(N_alphas)
for k in range(N_alphas):
print(k)
strehls[k] = calculate_average_strehl(alpha=alphas[k], coef=test_coef)
plt.figure()
plt.plot(alphas, strehls)
plt.xlabel(r'Alpha')
plt.ylabel(r'Strehl')
plt.show()
# Generate training and test datasets (clean PSF images)
train_PSF, train_coef, test_p, test_c = calibration.generate_dataset(PSF_zernike, N_train, N_test,
coef_strength=coef_strength, rescale=rescale)
# Check Defocus / Nominal ratio
peaks_nom = np.max(train_PSF[:, :, :, 0], axis=(1, 2))
peaks_foc = np.max(train_PSF[:, :, :, 1], axis=(1, 2))
plt.figure()
plt.plot(peaks_nom)
plt.plot(peaks_foc)
plt.show()
# Generate the Test Set using the coefficients
N_PSF = test_coef.shape[0]
test_PSF = np.zeros((N_PSF, pix, pix, 2))
for k in range(N_PSF):
pnom, _s = PSF_zernike.compute_PSF(test_coef[k])
pfoc, _s = PSF_zernike.compute_PSF(test_coef[k], diversity=True)
wave_r = waves_ratio[k]
psf_image, _strehl = PSFs.compute_PSF(c_act, wave_idx=k)
img = ax.imshow(psf_image, cmap=cmap, extent=[-1, 1, -1, 1])
ax.set_title('Wavelength: %.2f microns' % waves[k])
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
plt.colorbar(img, ax=ax, orientation='horizontal')
plt.show()
### ============================================================================================================ ###
# Generate the training sets
### ============================================================================================================ ###
# Generate a training set | calibration.generate_dataset automatically knows we are Multiwavelength
train_PSF, train_coef, test_PSF, test_coef = calibration.generate_dataset(PSFs, N_train, N_test,
coef_strength, rescale)
directory = os.path.join(os.getcwd(), 'Multiwave')
np.save(os.path.join(directory, 'train_PSF'), train_PSF)
np.save(os.path.join(directory, 'train_coef'), train_coef)
np.save(os.path.join(directory, 'test_PSF'), test_PSF)
np.save(os.path.join(directory, 'test_coef'), test_coef)
train_PSF = np.load(os.path.join(directory, 'train_PSF.npy'))
train_coef = np.load(os.path.join(directory, 'train_coef.npy'))
test_PSF = np.load(os.path.join(directory, 'test_PSF.npy'))
test_coef = np.load(os.path.join(directory, 'test_coef.npy'))
utils.show_PSF_multiwave(train_PSF)
plt.show()
train_coef = np.load(
os.path.join(directory, 'train_coef_%d.npy' % zern_level))
test_PSF = np.load(
os.path.join(directory, 'test_PSF_%d.npy' % zern_level))
test_coef = np.load(
os.path.join(directory, 'test_coef_%d.npy' % zern_level))
except:
# remember to rescale the coefficients. Otherwise the PSF images degrade more because of the fact that
# there are many more Zernikes
# we want to rescale the coefficients to ensure the minimum Strehl in the training set is sufficiently large
new_scale = rescale_coefficients_for_strehl(
PSF_zernike, coef_strength, rescale)
train_PSF, train_coef, test_PSF, test_coef = calibration.generate_dataset(
PSF_zernike, N_train, N_test, new_scale * coef_strength,
new_scale * rescale)
np.save(os.path.join(directory, 'train_PSF_%d' % zern_level),
train_PSF)
np.save(os.path.join(directory, 'train_coef_%d' % zern_level),
train_coef)
np.save(os.path.join(directory, 'test_PSF_%d' % zern_level),
test_PSF)
np.save(os.path.join(directory, 'test_coef_%d' % zern_level),
test_coef)
# Train a calibration model on the PSF images
calib = calibration.Calibration(PSF_model=PSF_zernike)
calib.create_cnn_model(layer_filters,
kernel_size,
defocus_zernike = np.zeros((1, zernike_matrix.shape[-1]))
defocus_zernike[0, 1] = 1.0
defocus_actuators = zernike_fit.fit_zernike_wave_to_actuators(defocus_zernike, plot=True, cmap='bwr')[:, 0]
# Update the Diversity Map on the actuator model so that it matches Defocus
diversity_defocus = diversity * defocus_actuators
PSF_actuators.define_diversity(diversity_defocus)
plt.figure()
plt.imshow(PSF_actuators.diversity_phase, cmap='RdBu')
plt.colorbar()
plt.title(r'Diversity Map | Defocus [rad]')
plt.show()
# Generate a training set for that nominal defocus
train_PSF, train_coef, test_PSF, test_coef = calibration.generate_dataset(PSF_actuators, N_train, N_test,
coef_strength, rescale)
utils.plot_images(train_PSF[5000:])
plt.show()
# Train the Calibration Model on images with the nominal defocus
calib = calibration.Calibration(PSF_model=PSF_actuators)
calib.create_cnn_model(layer_filers, kernel_size, name='CALIBR', activation='relu')
losses = calib.train_calibration_model(train_PSF, train_coef, test_PSF, test_coef,
N_loops, epochs_loop, verbose=1, batch_size_keras=32, plot_val_loss=False,
readout_noise=True, RMS_readout=[1. / SNR], readout_copies=readout_copies)
### Sometimes the train fails (no apparent reason) probably because of random weight initialization??
# If that happens, simply copy and paste the model definition and training bits, and try again
RMS_evolution, residual = calib.calibrate_iterations(test_PSF, test_coef, wavelength=WAVE, N_iter=N_iter,
readout_noise=True, RMS_readout=1./SNR)
readout_noise=True, RMS_readout=1./SNR)
calib.plot_RMS_evolution(RMS_evolution)
# plt.show()
### Up until here it's been the "classic" approach of building a model and training it
# ================================================================================================================ #
# Impact of Nyquist Errors
# ================================================================================================================ #
# Generate a set of PSF images with the wrong sampling
# Use the fake extra wavelength for which the sampling is off
PSF_error = psf.PointSpreadFunction(rbf_matrices[1], N_pix=N_PIX, crop_pix=pix, diversity_coef=diversity_defocus)
test_PSF_error, test_coef_error, _PSF, _coef = calibration.generate_dataset(PSF_error, 1000, 1,
coef_strength, rescale)
### WATCHOUT
# Calibrating here is quite tricky! Calib has a self.PSF_model the nominal model (not the one with Nyquist errors)
# that means that after the first iteration, the PSF images will revert to having the proper modelling
# We MUST temporarily override calib.PSF_model tem
calib.PSF_model = PSF_error
RMS_evolution_errors, residual_errors = calib.calibrate_iterations(test_PSF_error, test_coef_error,
wavelength=WAVE, N_iter=N_iter,
readout_noise=True, RMS_readout=1. / SNR)
calib.PSF_model = PSF_nom
calib.plot_RMS_evolution(RMS_evolution_errors)
# plt.show()