def _jwst_matrix_one_pair(norm, wfe_aber, resDir, savepsfs, saveopds, segment_pair):
"""
Function to calculate JWST mean contrast of one aberrated segment pair in NIRCam; for num_matrix_luvoir_multiprocess().
:param norm: float, direct PSF normalization factor (peak pixel of direct PSF)
:param wfe_aber: calibration aberration per segment in m
:param resDir: str, directory for matrix calculations
:param savepsfs: bool, if True, all PSFs will be saved to disk individually, as fits files
:param saveopds: bool, if True, all pupil surface maps of aberrated segment pairs will be saved to disk as PDF
:param segment_pair: tuple, pair of segments to aberrate, 0-indexed. If same segment gets passed in both tuple
entries, the segment will be aberrated only once.
Note how JWST segments start numbering at 0 just because that's python indexing, with 0 being
the segment A1.
:return: contrast as float, and segment pair as tuple
"""
# Set up JWST simulator in coronagraphic state
jwst_instrument, jwst_ote = webbpsf_imaging.set_up_nircam()
jwst_instrument.image_mask = CONFIG_PASTIS.get('JWST', 'focal_plane_mask')
# Put aberration on correct segments. If i=j, apply only once!
log.info(f'PAIR: {segment_pair[0]}-{segment_pair[1]}')
# Identify the correct JWST segments
seg_i = webbpsf_imaging.WSS_SEGS[segment_pair[0]].split('-')[0]
seg_j = webbpsf_imaging.WSS_SEGS[segment_pair[1]].split('-')[0]
# Put aberration on correct segments. If i=j, apply only once!
jwst_ote.zero()
jwst_ote.move_seg_local(seg_i, piston=wfe_aber, trans_unit='m')
if segment_pair[0] != segment_pair[1]:
jwst_ote.move_seg_local(seg_j, piston=wfe_aber, trans_unit='m')
log.info('Calculating coro image...')
image = jwst_instrument.calc_psf(nlambda=1)
psf = image[0].data / norm
# Save PSF image to disk
if savepsfs:
filename_psf = f'psf_piston_Noll1_segs_{segment_pair[0]}-{segment_pair[1]}'
hcipy.write_fits(psf, os.path.join(resDir, 'psfs', filename_psf + '.fits'))
# Plot segmented mirror WFE and save to disk
if saveopds:
opd_name = f'opd_piston_Noll1_segs_{segment_pair[0]}-{segment_pair[1]}'
plt.clf()
plt.figure(figsize=(8, 8))
ax2 = plt.subplot(111)
jwst_ote.display_opd(ax=ax2, vmax=500, colorbar_orientation='horizontal', title='Aberrated segment pair')
plt.savefig(os.path.join(resDir, 'OTE_images', opd_name + '.pdf'))
log.info('Calculating mean contrast in dark hole')
iwa = CONFIG_PASTIS.getfloat('JWST', 'IWA')
owa = CONFIG_PASTIS.getfloat('JWST', 'OWA')
sampling = CONFIG_PASTIS.getfloat('JWST', 'sampling')
dh_mask = util.create_dark_hole(psf, iwa, owa, sampling)
contrast = util.dh_mean(psf, dh_mask)
return contrast, segment_pair
def _hicat_matrix_one_pair(norm, wfe_aber, resDir, savepsfs, saveopds, segment_pair):
"""
Function to calculate HiCAT mean contrast of one aberrated segment pair; for num_matrix_luvoir_multiprocess().
:param norm: float, direct PSF normalization factor (peak pixel of direct PSF)
:param wfe_aber: calibration aberration per segment in m
:param resDir: str, directory for matrix calculations
:param savepsfs: bool, if True, all PSFs will be saved to disk individually, as fits files
:param saveopds: bool, if True, all pupil surface maps of aberrated segment pairs will be saved to disk as PDF
:param segment_pair: tuple, pair of segments to aberrate, 0-indexed. If same segment gets passed in both tuple
entries, the segment will be aberrated only once.
Note how HiCAT segments start numbering at 0, with 0 being the center segment.
:return: contrast as float, and segment pair as tuple
"""
# Set up HiCAT simulator in correct state
hicat_sim = set_up_hicat(apply_continuous_dm_maps=True)
hicat_sim.include_fpm = True
# Put aberration on correct segments. If i=j, apply only once!
log.info(f'PAIR: {segment_pair[0]}-{segment_pair[1]}')
hicat_sim.iris_dm.flatten()
hicat_sim.iris_dm.set_actuator(segment_pair[0], wfe_aber, 0, 0)
if segment_pair[0] != segment_pair[1]:
hicat_sim.iris_dm.set_actuator(segment_pair[1], wfe_aber, 0, 0)
log.info('Calculating coro image...')
image, inter = hicat_sim.calc_psf(display=False, return_intermediates=True)
psf = image[0].data / norm
# Save PSF image to disk
if savepsfs:
filename_psf = f'psf_piston_Noll1_segs_{segment_pair[0]}-{segment_pair[1]}'
hcipy.write_fits(psf, os.path.join(resDir, 'psfs', filename_psf + '.fits'))
# Plot segmented mirror WFE and save to disk
if saveopds:
opd_name = f'opd_piston_Noll1_segs_{segment_pair[0]}-{segment_pair[1]}'
plt.clf()
plt.imshow(inter[1].phase)
plt.savefig(os.path.join(resDir, 'OTE_images', opd_name + '.pdf'))
log.info('Calculating mean contrast in dark hole')
iwa = CONFIG_PASTIS.getfloat('HiCAT', 'IWA')
owa = CONFIG_PASTIS.getfloat('HiCAT', 'OWA')
sampling = CONFIG_PASTIS.getfloat('HiCAT', 'sampling')
dh_mask = util.create_dark_hole(psf, iwa, owa, sampling)
contrast = util.dh_mean(psf, dh_mask)
return contrast, segment_pair
def run_full_pastis_analysis(instrument,
run_choice,
design=None,
c_target=1e-10,
n_repeat=100):
"""
Run a full PASTIS analysis on a given PASTIS matrix.
The first couple of lines contain switches to turn different parts of the analysis on and off. These include:
1. calculating the PASTIS modes
2. calculating the PASTIS mode weights sigma under assumption of a uniform contrast allocation across all modes
3. running an E2E Monte Carlo simulation on the modes with their weights sigma from the uniform contrast allocation
4. calculating a cumulative contrast plot from the sigmas of the uniform contrast allocation
5. calculating the segment constraints mu under assumption of uniform statistical contrast contribution across segments
6. running an E2E Monte Carlo simulation on the segments with their weights mu
7. calculating the segment- and mode-space covariance matrices Ca and Cb
8. analytically calculating the statistical mean contrast and its variance
9. calculting segment-based error budget
:param instrument: str, "LUVOIR", "HiCAT" or "JWST"
:param run_choice: str, path to data and where outputs will be saved
:param design: str, optional, default=None, which means we read from the configfile (if running for LUVOIR):
what coronagraph design to use - 'small', 'medium' or 'large'
:param c_target: float, target contrast
:param n_repeat: number of realizations in both Monte Carlo simulations (modes and segments), default=100
"""
# Which parts are we running?
calculate_modes = True
calculate_sigmas = True
run_monte_carlo_modes = True
calc_cumulative_contrast = True
calculate_mus = True
run_monte_carlo_segments = True
calculate_covariance_matrices = True
analytical_statistics = True
calculate_segment_based = True
# Data directory
workdir = os.path.join(CONFIG_PASTIS.get('local', 'local_data_path'),
run_choice)
nseg = CONFIG_PASTIS.getint(instrument, 'nb_subapertures')
wvln = CONFIG_PASTIS.getfloat(instrument, 'lambda') * 1e-9 # [m]
log.info('Setting up optics...')
log.info(f'Data folder: {workdir}')
log.info(f'Instrument: {instrument}')
# Set up simulator, calculate reference PSF and dark hole mask
# TODO: replace this section with calculate_unaberrated_contrast_and_normalization(). This will require to save out
# reference and unaberrated coronagraphic PSF already in matrix generation.
if instrument == "LUVOIR":
if design is None:
design = CONFIG_PASTIS.get('LUVOIR', 'coronagraph_design')
log.info(f'Coronagraph design: {design}')
sampling = CONFIG_PASTIS.getfloat('LUVOIR', 'sampling')
optics_input = CONFIG_PASTIS.get('LUVOIR', 'optics_path')
luvoir = LuvoirAPLC(optics_input, design, sampling)
# Generate reference PSF and unaberrated coronagraphic image
luvoir.flatten()
psf_unaber, ref = luvoir.calc_psf(ref=True, display_intermediate=False)
norm = ref.max()
psf_unaber = psf_unaber.shaped / norm
dh_mask = luvoir.dh_mask.shaped
sim_instance = luvoir
if instrument == 'HiCAT':
hicat_sim = set_up_hicat(apply_continuous_dm_maps=True)
# Generate reference PSF and unaberrated coronagraphic image
hicat_sim.include_fpm = False
direct = hicat_sim.calc_psf()
norm = direct[0].data.max()
hicat_sim.include_fpm = True
coro_image = hicat_sim.calc_psf()
psf_unaber = coro_image[0].data / norm
# Create DH mask
iwa = CONFIG_PASTIS.getfloat('HiCAT', 'IWA')
owa = CONFIG_PASTIS.getfloat('HiCAT', 'OWA')
sampling = CONFIG_PASTIS.getfloat('HiCAT', 'sampling')
dh_mask = util.create_dark_hole(psf_unaber, iwa, owa,
sampling).astype('bool')
sim_instance = hicat_sim
if instrument == 'JWST':
jwst_sim = webbpsf_imaging.set_up_nircam(
) # this returns a tuple of two: jwst_sim[0] is the nircam object, jwst_sim[1] its ote
# Generate reference PSF and unaberrated coronagraphic image
jwst_sim[0].image_mask = None
direct = jwst_sim[0].calc_psf(nlambda=1)
direct_psf = direct[0].data
norm = direct_psf.max()
jwst_sim[0].image_mask = CONFIG_PASTIS.get('JWST', 'focal_plane_mask')
coro_image = jwst_sim[0].calc_psf(nlambda=1)
psf_unaber = coro_image[0].data / norm
# Create DH mask
iwa = CONFIG_PASTIS.getfloat('JWST', 'IWA')
owa = CONFIG_PASTIS.getfloat('JWST', 'OWA')
sampling = CONFIG_PASTIS.getfloat('JWST', 'sampling')
dh_mask = util.create_dark_hole(psf_unaber, iwa, owa,
sampling).astype('bool')
sim_instance = jwst_sim
# TODO: this would also be part of the refactor mentioned above
# Calculate coronagraph contrast floor
coro_floor = util.dh_mean(psf_unaber, dh_mask)
log.info(f'Coronagraph floor: {coro_floor}')
# Read the PASTIS matrix
matrix = fits.getdata(
os.path.join(workdir, 'matrix_numerical',
'PASTISmatrix_num_piston_Noll1.fits'))
### Calculate PASTIS modes and singular values/eigenvalues
if calculate_modes:
log.info('Calculating all PASTIS modes')
pmodes, svals = modes_from_matrix(instrument, workdir)
### Get full 2D modes and save them
mode_cube = full_modes_from_themselves(instrument,
pmodes,
workdir,
sim_instance,
saving=True)
else:
log.info(f'Reading PASTIS modes from {workdir}')
pmodes, svals = modes_from_file(workdir)
### Calculate mode-based static constraints
if calculate_sigmas:
log.info('Calculating static sigmas')
sigmas = calculate_sigma(c_target, nseg, svals, coro_floor)
np.savetxt(
os.path.join(workdir, 'results',
f'mode_requirements_{c_target}_uniform.txt'), sigmas)
# Plot static mode constraints
ppl.plot_mode_weights_simple(sigmas,
wvln,
out_dir=os.path.join(workdir, 'results'),
c_target=c_target,
fname_suffix='uniform',
save=True)
else:
log.info(f'Reading sigmas from {workdir}')
sigmas = np.loadtxt(
os.path.join(workdir, 'results',
f'mode_requirements_{c_target}_uniform.txt'))
### Calculate Monte Carlo simulation for sigmas, with E2E
if run_monte_carlo_modes:
log.info('\nRunning Monte Carlo simulation for modes')
# Keep track of time
start_monte_carlo_modes = time.time()
all_contr_rand_modes = []
all_random_weight_sets = []
for rep in range(n_repeat):
log.info(f'Mode realization {rep + 1}/{n_repeat}')
random_weights, one_contrast_mode = calc_random_mode_configurations(
instrument, pmodes, sim_instance, sigmas, dh_mask, norm)
all_random_weight_sets.append(random_weights)
all_contr_rand_modes.append(one_contrast_mode)
# Empirical mean and standard deviation of the distribution
mean_modes = np.mean(all_contr_rand_modes)
stddev_modes = np.std(all_contr_rand_modes)
log.info(f'Mean of the Monte Carlo result modes: {mean_modes}')
log.info(
f'Standard deviation of the Monte Carlo result modes: {stddev_modes}'
)
end_monte_carlo_modes = time.time()
# Save Monte Carlo simulation
np.savetxt(
os.path.join(workdir, 'results', f'mc_mode_reqs_{c_target}.txt'),
all_random_weight_sets)
np.savetxt(
os.path.join(workdir, 'results',
f'mc_modes_contrasts_{c_target}.txt'),
all_contr_rand_modes)
ppl.plot_monte_carlo_simulation(all_contr_rand_modes,
out_dir=os.path.join(
workdir, 'results'),
c_target=c_target,
segments=False,
stddev=stddev_modes,
save=True)
### Calculate cumulative contrast plot with E2E simulator and matrix product
if calc_cumulative_contrast:
log.info(
'Calculating cumulative contrast plot, uniform contrast across all modes'
)
cumulative_e2e = cumulative_contrast_e2e(instrument, pmodes, sigmas,
sim_instance, dh_mask, norm)
cumulative_pastis = cumulative_contrast_matrix(pmodes, sigmas, matrix,
coro_floor)
np.savetxt(
os.path.join(workdir, 'results',
f'cumul_contrast_accuracy_e2e_{c_target}.txt'),
cumulative_e2e)
np.savetxt(
os.path.join(workdir, 'results',
f'cumul_contrast_accuracy_pastis_{c_target}.txt'),
cumulative_pastis)
# Plot the cumulative contrast from E2E simulator and matrix
ppl.plot_cumulative_contrast_compare_accuracy(cumulative_pastis,
cumulative_e2e,
out_dir=os.path.join(
workdir, 'results'),
c_target=c_target,
save=True)
else:
log.info('Loading uniform cumulative contrast from disk.')
cumulative_e2e = np.loadtxt(
os.path.join(workdir, 'results',
f'cumul_contrast_accuracy_e2e_{c_target}.txt'))
### Calculate segment-based static constraints
if calculate_mus:
log.info('Calculating segment-based constraints')
mus = calculate_segment_constraints(pmodes, matrix, c_target,
coro_floor)
np.savetxt(
os.path.join(workdir, 'results',
f'segment_requirements_{c_target}.txt'), mus)
ppl.plot_segment_weights(mus,
out_dir=os.path.join(workdir, 'results'),
c_target=c_target,
save=True)
ppl.plot_mu_map(instrument,
mus,
sim_instance,
out_dir=os.path.join(workdir, 'results'),
c_target=c_target,
save=True)
# Apply mu map directly and run through E2E simulator
mus *= u.nm
if instrument == 'LUVOIR':
sim_instance.flatten()
for seg, mu in enumerate(mus):
sim_instance.set_segment(seg + 1, mu.to(u.m).value / 2, 0, 0)
im_data = sim_instance.calc_psf()
psf_pure_mu_map = im_data.shaped
if instrument == 'HiCAT':
sim_instance.iris_dm.flatten()
for seg, mu in enumerate(mus):
sim_instance.iris_dm.set_actuator(seg, mu / 1e9, 0,
0) # /1e9 converts to meters
im_data = sim_instance.calc_psf()
psf_pure_mu_map = im_data[0].data
if instrument == 'JWST':
sim_instance[1].zero()
for seg, mu in enumerate(mus):
seg_num = webbpsf_imaging.WSS_SEGS[seg].split('-')[0]
sim_instance[1].move_seg_local(seg_num,
piston=mu.value,
trans_unit='nm')
im_data = sim_instance[0].calc_psf(nlambda=1)
psf_pure_mu_map = im_data[0].data
contrast_mu = util.dh_mean(psf_pure_mu_map / norm, dh_mask)
log.info(f'Contrast with pure mu-map: {contrast_mu}')
else:
log.info(f'Reading mus from {workdir}')
mus = np.loadtxt(
os.path.join(workdir, 'results',
f'segment_requirements_{c_target}.txt'))
mus *= u.nm
### Calculate Monte Carlo confirmation for segments, with E2E
if run_monte_carlo_segments:
log.info('\nRunning Monte Carlo simulation for segments')
# Keep track of time
start_monte_carlo_seg = time.time()
all_contr_rand_seg = []
all_random_maps = []
for rep in range(n_repeat):
log.info(f'Segment realization {rep + 1}/{n_repeat}')
random_map, one_contrast_seg = calc_random_segment_configuration(
instrument, sim_instance, mus, dh_mask, norm)
all_random_maps.append(random_map)
all_contr_rand_seg.append(one_contrast_seg)
# Empirical mean and standard deviation of the distribution
mean_segments = np.mean(all_contr_rand_seg)
stddev_segments = np.std(all_contr_rand_seg)
log.info(f'Mean of the Monte Carlo result segments: {mean_segments}')
log.info(
f'Standard deviation of the Monte Carlo result segments: {stddev_segments}'
)
with open(
os.path.join(workdir, 'results',
f'statistical_contrast_empirical_{c_target}.txt'),
'w') as file:
file.write(f'Empirical, statistical mean: {mean_segments}')
file.write(f'\nEmpirical variance: {stddev_segments**2}')
end_monte_carlo_seg = time.time()
log.info('\nRuntimes:')
log.info(
'Monte Carlo on segments with {} iterations: {} sec = {} min = {} h'
.format(n_repeat, end_monte_carlo_seg - start_monte_carlo_seg,
(end_monte_carlo_seg - start_monte_carlo_seg) / 60,
(end_monte_carlo_seg - start_monte_carlo_seg) / 3600))
# Save Monte Carlo simulation
np.savetxt(
os.path.join(workdir, 'results',
f'mc_segment_req_maps_{c_target}.txt'),
all_random_maps) # in m
np.savetxt(
os.path.join(workdir, 'results',
f'mc_segments_contrasts_{c_target}.txt'),
all_contr_rand_seg)
ppl.plot_monte_carlo_simulation(all_contr_rand_seg,
out_dir=os.path.join(
workdir, 'results'),
c_target=c_target,
segments=True,
stddev=stddev_segments,
save=True)
### Calculate covariance matrices
if calculate_covariance_matrices:
log.info('Calculating covariance matrices')
Ca = np.diag(np.square(mus.value))
hcipy.write_fits(
Ca,
os.path.join(
workdir, 'results',
f'cov_matrix_segments_Ca_{c_target}_segment-based.fits'))
Cb = np.dot(np.transpose(pmodes), np.dot(Ca, pmodes))
hcipy.write_fits(
Cb,
os.path.join(workdir, 'results',
f'cov_matrix_modes_Cb_{c_target}_segment-based.fits'))
ppl.plot_covariance_matrix(Ca,
os.path.join(workdir, 'results'),
c_target,
segment_space=True,
fname_suffix='segment-based',
save=True)
ppl.plot_covariance_matrix(Cb,
os.path.join(workdir, 'results'),
c_target,
segment_space=False,
fname_suffix='segment-based',
save=True)
else:
log.info('Loading covariance matrices from disk.')
Ca = fits.getdata(
os.path.join(
workdir, 'results',
f'cov_matrix_segments_Ca_{c_target}_segment-based.fits'))
Cb = fits.getdata(
os.path.join(workdir, 'results',
f'cov_matrix_modes_Cb_{c_target}_segment-based.fits'))
### Analytically calculate statistical mean contrast and its variance
if analytical_statistics:
log.info('Calculating analytical statistics.')
mean_stat_c = util.calc_statistical_mean_contrast(
matrix, Ca, coro_floor)
var_c = util.calc_variance_of_mean_contrast(matrix, Ca)
log.info(f'Analytical statistical mean: {mean_stat_c}')
log.info(f'Analytical standard deviation: {np.sqrt(var_c)}')
with open(
os.path.join(
workdir, 'results',
f'statistical_contrast_analytical_{c_target}.txt'),
'w') as file:
file.write(f'Analytical, statistical mean: {mean_stat_c}')
file.write(f'\nAnalytical variance: {var_c}')
### Calculate segment-based error budget
if calculate_segment_based:
log.info('Calculating segment-based error budget.')
# Extract segment-based mode weights
log.info('Calculate segment-based mode weights')
sigmas_opt = np.sqrt(np.diag(Cb))
np.savetxt(
os.path.join(workdir, 'results',
f'mode_requirements_{c_target}_segment-based.txt'),
sigmas_opt)
ppl.plot_mode_weights_simple(sigmas_opt,
wvln,
out_dir=os.path.join(workdir, 'results'),
c_target=c_target,
fname_suffix='segment-based',
save=True)
ppl.plot_mode_weights_double_axis(
(sigmas, sigmas_opt),
wvln,
os.path.join(workdir, 'results'),
c_target,
fname_suffix='segment-based-vs-uniform',
labels=('Uniform error budget', 'Segment-based error budget'),
alphas=(0.5, 1.),
linestyles=('--', '-'),
colors=('k', 'r'),
save=True)
# Calculate contrast per mode
log.info('Calculating contrast per mode')
per_mode_opt_e2e = cumulative_contrast_e2e(instrument,
pmodes,
sigmas_opt,
sim_instance,
dh_mask,
norm,
individual=True)
np.savetxt(
os.path.join(
workdir, 'results',
f'contrast_per_mode_{c_target}_e2e_segment-based.txt'),
per_mode_opt_e2e)
ppl.plot_contrast_per_mode(per_mode_opt_e2e,
coro_floor,
c_target,
pmodes.shape[0],
os.path.join(workdir, 'results'),
save=True)
# Calculate segment-based cumulative contrast
log.info('Calculating segment-based cumulative contrast')
cumulative_opt_e2e = cumulative_contrast_e2e(instrument, pmodes,
sigmas_opt, sim_instance,
dh_mask, norm)
np.savetxt(
os.path.join(
workdir, 'results',
f'cumul_contrast_allocation_e2e_{c_target}_segment-based.txt'),
cumulative_opt_e2e)
# Plot cumulative contrast from E2E simulator, segment-based vs. uniform error budget
ppl.plot_cumulative_contrast_compare_allocation(
cumulative_opt_e2e,
cumulative_e2e,
os.path.join(workdir, 'results'),
c_target,
fname_suffix='segment-based-vs-uniform',
save=True)
### Write full PDF report
title_page_list = util.collect_title_page(workdir, c_target)
util.create_title_page(instrument, workdir, title_page_list)
util.create_pdf_report(workdir, c_target)
### DONE
log.info(f"All saved in {os.path.join(workdir, 'results')}")
log.info('Good job')
def calculate_unaberrated_contrast_and_normalization(instrument, design=None, return_coro_simulator=True, save_coro_floor=False, save_psfs=False, outpath=''):
"""
Calculate the direct PSF peak and unaberrated coronagraph floor of an instrument.
:param instrument: string, 'LUVOIR', 'HiCAT' or 'JWST'
:param design: str, optional, default=None, which means we read from the configfile: what coronagraph design
to use - 'small', 'medium' or 'large'
:param return_coro_simulator: bool, whether to return the coronagraphic simulator as third return, default True
:param save: bool, if True, will save direct and coro PSF images to disk, default False
:param outpath: string, where to save outputs to if save=True
:return: contrast floor and PSF normalization factor, and optionally (by default) the simulator in coron mode
"""
if instrument == 'LUVOIR':
# Instantiate LuvoirAPLC class
sampling = CONFIG_PASTIS.getfloat(instrument, 'sampling')
optics_input = CONFIG_PASTIS.get('LUVOIR', 'optics_path')
if design is None:
design = CONFIG_PASTIS.get('LUVOIR', 'coronagraph_design')
luvoir = LuvoirAPLC(optics_input, design, sampling)
# Calculate reference images for contrast normalization and coronagraph floor
unaberrated_coro_psf, direct = luvoir.calc_psf(ref=True, display_intermediate=False, return_intermediate=False)
norm = np.max(direct)
direct_psf = direct.shaped
coro_psf = unaberrated_coro_psf.shaped / norm
# Return the coronagraphic simulator and DH mask
coro_simulator = luvoir
dh_mask = luvoir.dh_mask.shaped
if instrument == 'HiCAT':
# Set up HiCAT simulator in correct state
hicat_sim = set_up_hicat(apply_continuous_dm_maps=True)
# Calculate direct reference images for contrast normalization
hicat_sim.include_fpm = False
direct = hicat_sim.calc_psf()
direct_psf = direct[0].data
norm = direct_psf.max()
# Calculate unaberrated coronagraph image for contrast floor
hicat_sim.include_fpm = True
coro_image = hicat_sim.calc_psf()
coro_psf = coro_image[0].data / norm
iwa = CONFIG_PASTIS.getfloat('HiCAT', 'IWA')
owa = CONFIG_PASTIS.getfloat('HiCAT', 'OWA')
sampling = CONFIG_PASTIS.getfloat('HiCAT', 'sampling')
dh_mask = util.create_dark_hole(coro_psf, iwa, owa, sampling).astype('bool')
# Return the coronagraphic simulator
coro_simulator = hicat_sim
if instrument == 'JWST':
# Instantiate NIRCAM object
jwst_sim = webbpsf_imaging.set_up_nircam() # this returns a tuple of two: jwst_sim[0] is the nircam object, jwst_sim[1] its ote
# Calculate direct reference images for contrast normalization
jwst_sim[0].image_mask = None
direct = jwst_sim[0].calc_psf(nlambda=1)
direct_psf = direct[0].data
norm = direct_psf.max()
# Calculate unaberrated coronagraph image for contrast floor
jwst_sim[0].image_mask = CONFIG_PASTIS.get('JWST', 'focal_plane_mask')
coro_image = jwst_sim[0].calc_psf(nlambda=1)
coro_psf = coro_image[0].data / norm
iwa = CONFIG_PASTIS.getfloat('JWST', 'IWA')
owa = CONFIG_PASTIS.getfloat('JWST', 'OWA')
sampling = CONFIG_PASTIS.getfloat('JWST', 'sampling')
dh_mask = util.create_dark_hole(coro_psf, iwa, owa, sampling).astype('bool')
# Return the coronagraphic simulator (a tuple in the JWST case!)
coro_simulator = jwst_sim
# Calculate coronagraph floor in dark hole
contrast_floor = util.dh_mean(coro_psf, dh_mask)
log.info(f'contrast floor: {contrast_floor}')
if save_coro_floor:
# Save contrast floor to text file
with open(os.path.join(outpath, 'coronagraph_floor.txt'), 'w') as file:
file.write(f'Coronagraph floor: {contrast_floor}')
if save_psfs:
# Save direct PSF, unaberrated coro PSF and DH masked coro PSF as PDF
plt.figure(figsize=(18, 6))
plt.subplot(1, 3, 1)
plt.title("Direct PSF")
plt.imshow(direct_psf, norm=LogNorm())
plt.colorbar()
plt.subplot(1, 3, 2)
plt.title("Unaberrated coro PSF")
plt.imshow(coro_psf, norm=LogNorm())
plt.colorbar()
plt.subplot(1, 3, 3)
plt.title("Dark hole coro PSF")
plt.imshow(np.ma.masked_where(~dh_mask, coro_psf), norm=LogNorm())
plt.colorbar()
plt.savefig(os.path.join(outpath, 'unaberrated_dh.pdf'))
if return_coro_simulator:
return contrast_floor, norm, coro_simulator
else:
return contrast_floor, norm
def num_matrix_jwst():
"""
Generate a numerical PASTIS matrix for a JWST coronagraph.
-- Depracated function, the LUVOIR PASTIS matrix is better calculated with num_matrix_multiprocess(), which can
do this for your choice of one of the implemented instruments (LUVOIR, HiCAT, JWST). --
All inputs are read from the (local) configfile and saved to the specified output directory.
"""
import webbpsf
from e2e_simulators import webbpsf_imaging as webbim
# Set WebbPSF environment variable
os.environ['WEBBPSF_PATH'] = CONFIG_PASTIS.get('local', 'webbpsf_data_path')
# Keep track of time
start_time = time.time() # runtime is currently around 21 minutes
log.info('Building numerical matrix for JWST\n')
# Parameters
overall_dir = util.create_data_path(CONFIG_PASTIS.get('local', 'local_data_path'), telescope='jwst')
resDir = os.path.join(overall_dir, 'matrix_numerical')
which_tel = CONFIG_PASTIS.get('telescope', 'name')
nb_seg = CONFIG_PASTIS.getint(which_tel, 'nb_subapertures')
im_size_e2e = CONFIG_PASTIS.getint('numerical', 'im_size_px_webbpsf')
inner_wa = CONFIG_PASTIS.getint(which_tel, 'IWA')
outer_wa = CONFIG_PASTIS.getint(which_tel, 'OWA')
sampling = CONFIG_PASTIS.getfloat(which_tel, 'sampling')
fpm = CONFIG_PASTIS.get(which_tel, 'focal_plane_mask') # focal plane mask
lyot_stop = CONFIG_PASTIS.get(which_tel, 'pupil_plane_stop') # Lyot stop
filter = CONFIG_PASTIS.get(which_tel, 'filter_name')
wfe_aber = CONFIG_PASTIS.getfloat(which_tel, 'calibration_aberration') * u.nm
wss_segs = webbpsf.constants.SEGNAMES_WSS_ORDER
zern_max = CONFIG_PASTIS.getint('zernikes', 'max_zern')
zern_number = CONFIG_PASTIS.getint('calibration', 'local_zernike')
zern_mode = util.ZernikeMode(zern_number) # Create Zernike mode object for easier handling
wss_zern_nb = util.noll_to_wss(zern_number) # Convert from Noll to WSS framework
# Create necessary directories if they don't exist yet
os.makedirs(overall_dir, exist_ok=True)
os.makedirs(resDir, exist_ok=True)
os.makedirs(os.path.join(resDir, 'OTE_images'), exist_ok=True)
os.makedirs(os.path.join(resDir, 'psfs'), exist_ok=True)
os.makedirs(os.path.join(resDir, 'darkholes'), exist_ok=True)
# Create the dark hole mask.
pup_im = np.zeros([im_size_e2e, im_size_e2e]) # this is just used for DH mask generation
dh_area = util.create_dark_hole(pup_im, inner_wa, outer_wa, sampling)
# Create a direct WebbPSF image for normalization factor
fake_aber = np.zeros([nb_seg, zern_max])
psf_perfect = webbim.nircam_nocoro(filter, fake_aber)
normp = np.max(psf_perfect)
psf_perfect = psf_perfect / normp
# Set up NIRCam coro object from WebbPSF
nc_coro = webbpsf.NIRCam()
nc_coro.filter = filter
nc_coro.image_mask = fpm
nc_coro.pupil_mask = lyot_stop
# Null the OTE OPDs for the PSFs, maybe we will add internal WFE later.
nc_coro, ote_coro = webbpsf.enable_adjustable_ote(nc_coro) # create OTE for coronagraph
nc_coro.include_si_wfe = False # set SI internal WFE to zero
#-# Generating the PASTIS matrix and a list for all contrasts
contrast_matrix = np.zeros([nb_seg, nb_seg]) # Generate empty matrix
all_psfs = []
all_dhs = []
all_contrasts = []
log.info(f'wfe_aber: {wfe_aber}')
for i in range(nb_seg):
for j in range(nb_seg):
log.info(f'\nSTEP: {i+1}-{j+1} / {nb_seg}-{nb_seg}')
# Get names of segments, they're being addressed by their names in the ote functions.
seg_i = wss_segs[i].split('-')[0]
seg_j = wss_segs[j].split('-')[0]
# Put the aberration on the correct segments
Aber_WSS = np.zeros([nb_seg, zern_max]) # The Zernikes here will be filled in the WSS order!!!
# Because it goes into _apply_hexikes_to_seg().
Aber_WSS[i, wss_zern_nb - 1] = wfe_aber.to(u.m).value # Aberration on the segment we're currently working on;
# convert to meters; -1 on the Zernike because Python starts
# numbering at 0.
Aber_WSS[j, wss_zern_nb - 1] = wfe_aber.to(u.m).value # same for other segment
# Putting aberrations on segments i and j
ote_coro.reset() # Making sure there are no previous movements on the segments.
ote_coro.zero() # set OTE for coronagraph to zero
# Apply both aberrations to OTE. If i=j, apply only once!
ote_coro._apply_hexikes_to_seg(seg_i, Aber_WSS[i, :]) # set segment i (segment numbering starts at 1)
if i != j:
ote_coro._apply_hexikes_to_seg(seg_j, Aber_WSS[j, :]) # set segment j
# If you want to display it:
# ote_coro.display_opd()
# plt.show()
# Save OPD images for testing
opd_name = f'opd_{zern_mode.name}_{zern_mode.convention + str(zern_mode.index)}_segs_{i+1}-{j+1}'
plt.clf()
ote_coro.display_opd()
plt.savefig(os.path.join(resDir, 'OTE_images', opd_name + '.pdf'))
log.info('Calculating WebbPSF image')
image = nc_coro.calc_psf(fov_pixels=int(im_size_e2e), oversample=1, nlambda=1)
psf = image[0].data / normp
# Save WebbPSF image to disk
filename_psf = f'psf_{zern_mode.name}_{zern_mode.convention + str(zern_mode.index)}_segs_{i+1}-{j+1}'
util.write_fits(psf, os.path.join(resDir, 'psfs', filename_psf + '.fits'), header=None, metadata=None)
all_psfs.append(psf)
log.info('Calculating mean contrast in dark hole')
dh_intensity = psf * dh_area
contrast = np.mean(dh_intensity[np.where(dh_intensity != 0)])
log.info(f'contrast: {contrast}')
# Save DH image to disk and put current contrast in list
filename_dh = f'dh_{zern_mode.name}_{zern_mode.convention + str(zern_mode.index)}_segs_{i+1}-{j+1}'
util.write_fits(dh_intensity, os.path.join(resDir, 'darkholes', filename_dh + '.fits'), header=None, metadata=None)
all_dhs.append(dh_intensity)
all_contrasts.append(contrast)
# Fill according entry in the matrix
contrast_matrix[i,j] = contrast
# Transform saved lists to arrays
all_psfs = np.array(all_psfs)
all_dhs = np.array(all_dhs)
all_contrasts = np.array(all_contrasts)
# Filling the off-axis elements
matrix_two_N = np.copy(contrast_matrix) # This is just an intermediary copy so that I don't mix things up.
matrix_pastis = np.copy(contrast_matrix) # This will be the final PASTIS matrix.
for i in range(nb_seg):
for j in range(nb_seg):
if i != j:
matrix_off_val = (matrix_two_N[i,j] - matrix_two_N[i,i] - matrix_two_N[j,j]) / 2.
matrix_pastis[i,j] = matrix_off_val
log.info(f'Off-axis for i{i+1}-j{j+1}: {matrix_off_val}')
# Normalize matrix for the input aberration
matrix_pastis /= np.square(wfe_aber.value)
# Save matrix to file
filename_matrix = f'PASTISmatrix_num_{zern_mode.name}_{zern_mode.convention + str(zern_mode.index)}'
util.write_fits(matrix_pastis, os.path.join(resDir, filename_matrix + '.fits'), header=None, metadata=None)
log.info(f'Matrix saved to: {os.path.join(resDir, filename_matrix + ".fits")}')
# Save the PSF and DH image *cubes* as well (as opposed to each one individually)
util.write_fits(all_psfs, os.path.join(resDir, 'psfs', 'psf_cube.fits'), header=None, metadata=None)
util.write_fits(all_dhs, os.path.join(resDir, 'darkholes', 'dh_cube.fits'), header=None, metadata=None)
np.savetxt(os.path.join(resDir, 'pair-wise_contrasts.txt'), all_contrasts, fmt='%e')
# Tell us how long it took to finish.
end_time = time.time()
log.info(f'Runtime for matrix_building.py: {end_time - start_time}sec = {(end_time - start_time) / 60}min')
log.info(f'Data saved to {resDir}')
normp = np.max(psf_default)
psf_default = psf_default / normp
psf_coro = psf_coro / normp
# Save the PSFs for testing
util.write_fits(psf_default,
os.path.join(outDir, 'psf_default.fits'),
header=None,
metadata=None)
util.write_fits(psf_coro,
os.path.join(outDir, 'psf_coro.fits'),
header=None,
metadata=None)
# Create the dark hole
dh_area = util.create_dark_hole(psf_coro, inner_wa, outer_wa, sampling)
util.write_fits(dh_area,
os.path.join(outDir, 'dh_area.fits'),
header=None,
metadata=None)
# Calculate the baseline contrast *with* the coronagraph and *without* aberrations and save the value to file
contrast_im = psf_coro * dh_area
contrast_base = np.mean(contrast_im[np.where(contrast_im != 0)])
contrastname = 'base-contrast_' + zern_mode.name + '_' + zern_mode.convention + str(
zern_mode.index
) #TODO: Why does the filename include a Zernike if this is supposed to be the perfect PSF without aberrations?
contrast_fake_array = np.array(contrast_base).reshape(
1,
) # Convert into array of shape (1,), otherwise np.savetxt() doesn't work
np.savetxt(os.path.join(outDir, contrastname + '.txt'),
def analytical_model(zernike_pol, coef, cali=False):
"""
:param zernike_pol:
:param coef:
:param cali: bool; True if we already have calibration coefficients to use. False if we still need to create them.
:return:
"""
#-# Parameters
dataDir = os.path.join(CONFIG_PASTIS.get('local', 'local_data_path'),
'active')
telescope = CONFIG_PASTIS.get('telescope', 'name')
nb_seg = CONFIG_PASTIS.getint(telescope, 'nb_subapertures')
tel_size_m = CONFIG_PASTIS.getfloat(telescope, 'diameter') * u.m
real_size_seg = CONFIG_PASTIS.getfloat(
telescope, 'flat_to_flat'
) # in m, size in meters of an individual segment flatl to flat
size_seg = CONFIG_PASTIS.getint(
'numerical',
'size_seg') # pixel size of an individual segment tip to tip
wvln = CONFIG_PASTIS.getint(telescope, 'lambda') * u.nm
inner_wa = CONFIG_PASTIS.getint(telescope, 'IWA')
outer_wa = CONFIG_PASTIS.getint(telescope, 'OWA')
tel_size_px = CONFIG_PASTIS.getint(
'numerical', 'tel_size_px') # pupil diameter of telescope in pixels
im_size_pastis = CONFIG_PASTIS.getint(
'numerical', 'im_size_px_pastis') # image array size in px
sampling = CONFIG_PASTIS.getfloat(telescope, 'sampling') # sampling
size_px_tel = tel_size_m / tel_size_px # size of one pixel in pupil plane in m
px_sq_to_rad = (size_px_tel * np.pi / tel_size_m) * u.rad
zern_max = CONFIG_PASTIS.getint('zernikes', 'max_zern')
sz = CONFIG_PASTIS.getint(
'ATLAST',
'im_size_lamD_hcipy') # image size in lam/D, only used in ATLAST case
# Create Zernike mode object for easier handling
zern_mode = util.ZernikeMode(zernike_pol)
#-# Mean subtraction for piston
if zernike_pol == 1:
coef -= np.mean(coef)
#-# Generic segment shapes
if telescope == 'JWST':
# Load pupil from file
pupil = fits.getdata(
os.path.join(dataDir, 'segmentation', 'pupil.fits'))
# Put pupil in randomly picked, slightly larger image array
pup_im = np.copy(pupil) # remove if lines below this are active
#pup_im = np.zeros([tel_size_px, tel_size_px])
#lim = int((pup_im.shape[1] - pupil.shape[1])/2.)
#pup_im[lim:-lim, lim:-lim] = pupil
# test_seg = pupil[394:,197:315] # this is just so that I can display an individual segment when the pupil is 512
# test_seg = pupil[:203,392:631] # ... when the pupil is 1024
# one_seg = np.zeros_like(test_seg)
# one_seg[:110, :] = test_seg[8:, :] # this is the centered version of the individual segment for 512 px pupil
# Creat a mini-segment (one individual segment from the segmented aperture)
mini_seg_real = poppy.NgonAperture(
name='mini', radius=real_size_seg
) # creating real mini segment shape with poppy
#test = mini_seg_real.sample(wavelength=wvln, grid_size=flat_diam, return_scale=True) # fix its sampling with wavelength
mini_hdu = mini_seg_real.to_fits(wavelength=wvln,
npix=size_seg) # make it a fits file
mini_seg = mini_hdu[
0].data # extract the image data from the fits file
elif telescope == 'ATLAST':
# Create mini-segment
pupil_grid = hcipy.make_pupil_grid(dims=tel_size_px,
diameter=real_size_seg)
focal_grid = hcipy.make_focal_grid(
pupil_grid, sampling, sz, wavelength=wvln.to(
u.m).value) # fov = lambda/D radius of total image
prop = hcipy.FraunhoferPropagator(pupil_grid, focal_grid)
mini_seg_real = hcipy.hexagonal_aperture(circum_diameter=real_size_seg,
angle=np.pi / 2)
mini_seg_hc = hcipy.evaluate_supersampled(
mini_seg_real, pupil_grid, 4
) # the supersampling number doesn't really matter in context with the other numbers
mini_seg = mini_seg_hc.shaped # make it a 2D array
# Redefine size_seg if using HCIPy
size_seg = mini_seg.shape[0]
# Make stand-in pupil for DH array
pupil = fits.getdata(
os.path.join(dataDir, 'segmentation', 'pupil.fits'))
pup_im = np.copy(pupil)
#-# Generate a dark hole mask
#TODO: simplify DH generation and usage
dh_area = util.create_dark_hole(
pup_im, inner_wa, outer_wa, sampling
) # this might become a problem if pupil size is not same like pastis image size. fine for now though.
if telescope == 'ATLAST':
dh_sz = util.zoom_cen(dh_area, sz * sampling)
#-# Import information form segmentation script
Projection_Matrix = fits.getdata(
os.path.join(dataDir, 'segmentation', 'Projection_Matrix.fits'))
vec_list = fits.getdata(
os.path.join(dataDir, 'segmentation', 'vec_list.fits')) # in pixels
NR_pairs_list = fits.getdata(
os.path.join(dataDir, 'segmentation', 'NR_pairs_list_int.fits'))
# Figure out how many NRPs we're dealing with
NR_pairs_nb = NR_pairs_list.shape[0]
#-# Chose whether calibration factors to do the calibraiton with
if cali:
filename = 'calibration_' + zern_mode.name + '_' + zern_mode.convention + str(
zern_mode.index)
ck = fits.getdata(
os.path.join(dataDir, 'calibration', filename + '.fits'))
else:
ck = np.ones(nb_seg)
coef = coef * ck
#-# Generic coefficients
# the coefficients in front of the non redundant pairs, the A_q in eq. 13 in Leboulleux et al. 2018
generic_coef = np.zeros(
NR_pairs_nb
) * u.nm * u.nm # setting it up with the correct units this will have
for q in range(NR_pairs_nb):
for i in range(nb_seg):
for j in range(i + 1, nb_seg):
if Projection_Matrix[i, j, 0] == q + 1:
generic_coef[q] += coef[i] * coef[j]
#-# Constant sum and cosine sum - calculating eq. 13 from Leboulleux et al. 2018
if telescope == 'JWST':
i_line = np.linspace(-im_size_pastis / 2., im_size_pastis / 2.,
im_size_pastis)
tab_i, tab_j = np.meshgrid(i_line, i_line)
cos_u_mat = np.zeros(
(int(im_size_pastis), int(im_size_pastis), NR_pairs_nb))
elif telescope == 'ATLAST':
i_line = np.linspace(-(2 * sz * sampling) / 2.,
(2 * sz * sampling) / 2., (2 * sz * sampling))
tab_i, tab_j = np.meshgrid(i_line, i_line)
cos_u_mat = np.zeros((int((2 * sz * sampling)), int(
(2 * sz * sampling)), NR_pairs_nb))
# Calculating the cosine terms from eq. 13.
# The -1 with each NR_pairs_list is because the segment names are saved starting from 1, but Python starts
# its indexing at zero, so we have to make it start at zero here too.
for q in range(NR_pairs_nb):
# cos(b_q <dot> u): b_q with 1 <= q <= NR_pairs_nb is the basis of NRPS, meaning the distance vectors between
# two segments of one NRP. We can read these out from vec_list.
# u is the position (vector) in the detector plane. Here, those are the grids tab_i and tab_j.
# We need to calculate the dot product between all b_q and u, so in each iteration (for q), we simply add the
# x and y component.
cos_u_mat[:, :, q] = np.cos(
px_sq_to_rad *
(vec_list[NR_pairs_list[q, 0] - 1, NR_pairs_list[q, 1] - 1, 0] *
tab_i) + px_sq_to_rad *
(vec_list[NR_pairs_list[q, 0] - 1, NR_pairs_list[q, 1] - 1, 1] *
tab_j)) * u.dimensionless_unscaled
sum1 = np.sum(
coef**2
) # sum of all a_{k,l} in eq. 13 - this works only for single Zernikes (l fixed), because np.sum would sum over l too, which would be wrong.
if telescope == 'JWST':
sum2 = np.zeros(
(int(im_size_pastis), int(im_size_pastis))
) * u.nm * u.nm # setting it up with the correct units this will have
elif telescope == 'ATLAST':
sum2 = np.zeros(
(int(2 * sz * sampling), int(2 * sz * sampling))) * u.nm * u.nm
for q in range(NR_pairs_nb):
sum2 = sum2 + generic_coef[q] * cos_u_mat[:, :, q]
#-# Local Zernike
if telescope == 'JWST':
# Generate a basis of Zernikes with the mini segment being the support
isolated_zerns = zern.hexike_basis(nterms=zern_max,
npix=size_seg,
rho=None,
theta=None,
vertical=False,
outside=0.0)
# Calculate the Zernike that is currently being used and put it on one single subaperture, the result is Zer
# Apply the currently used Zernike to the mini-segment.
if zernike_pol == 1:
Zer = np.copy(mini_seg)
elif zernike_pol in range(2, zern_max - 2):
Zer = np.copy(mini_seg)
Zer = Zer * isolated_zerns[zernike_pol - 1]
# Fourier Transform of the Zernike - the global envelope
mf = mft.MatrixFourierTransform()
ft_zern = mf.perform(Zer, im_size_pastis / sampling, im_size_pastis)
elif telescope == 'ATLAST':
isolated_zerns = hcipy.make_zernike_basis(num_modes=zern_max,
D=real_size_seg,
grid=pupil_grid,
radial_cutoff=False)
Zer = hcipy.Wavefront(mini_seg_hc * isolated_zerns[zernike_pol - 1],
wavelength=wvln.to(u.m).value)
# Fourier transform the Zernike
ft_zern = prop(Zer)
#-# Final image
if telescope == 'JWST':
# Generating the final image that will get passed on to the outer scope, I(u) in eq. 13
intensity = np.abs(ft_zern)**2 * (sum1.value + 2. * sum2.value)
elif telescope == 'ATLAST':
intensity = ft_zern.intensity.shaped * (sum1.value + 2. * sum2.value)
# PASTIS is only valid inside the dark hole, so we cut out only that part
if telescope == 'JWST':
tot_dh_im_size = sampling * (outer_wa + 3)
intensity_zoom = util.zoom_cen(
intensity, tot_dh_im_size
) # zoom box is (owa + 3*lambda/D) wide, in terms of lambda/D
dh_area_zoom = util.zoom_cen(dh_area, tot_dh_im_size)
dh_psf = dh_area_zoom * intensity_zoom
elif telescope == 'ATLAST':
dh_psf = dh_sz * intensity
"""
# Create plots.
plt.subplot(1, 3, 1)
plt.imshow(pupil, origin='lower')
plt.title('JWST pupil and diameter definition')
plt.plot([46.5, 464.5], [101.5, 409.5], 'r-') # show how the diagonal of the pupil is defined
plt.subplot(1, 3, 2)
plt.imshow(mini_seg, origin='lower')
plt.title('JWST individual mini-segment')
plt.subplot(1, 3, 3)
plt.imshow(dh_psf, origin='lower')
plt.title('JWST dark hole')
plt.show()
"""
# dh_psf is the image of the dark hole only, the pixels outside of it are zero
# intensity is the entire final image
return dh_psf, intensity