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 _luvoir_matrix_one_pair(design, norm, wfe_aber, zern_mode, resDir, savepsfs, saveopds, segment_pair):
"""
Function to calculate LVUOIR-A mean contrast of one aberrated segment pair; for num_matrix_luvoir_multiprocess().
:param design: str, what coronagraph design to use - 'small', 'medium' or 'large'
:param norm: float, direct PSF normalization factor (peak pixel of direct PSF)
:param wfe_aber: float, calibration aberration per segment in m
:param zern_mode: Zernike mode object, local Zernike aberration
: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 LUVOIR segments start numbering at 1, with 0 being the center segment that doesn't exist.
:return: contrast as float, and segment pair as tuple
"""
# Instantiate LUVOIR object
sampling = CONFIG_PASTIS.getfloat('LUVOIR', 'sampling')
optics_input = CONFIG_PASTIS.get('LUVOIR', 'optics_path')
luv = LuvoirAPLC(optics_input, design, sampling)
log.info(f'PAIR: {segment_pair[0]+1}-{segment_pair[1]+1}')
# Put aberration on correct segments. If i=j, apply only once!
luv.flatten()
luv.set_segment(segment_pair[0]+1, wfe_aber / 2, 0, 0)
if segment_pair[0] != segment_pair[1]:
luv.set_segment(segment_pair[1]+1, wfe_aber / 2, 0, 0)
log.info('Calculating coro image...')
image, inter = luv.calc_psf(ref=False, display_intermediate=False, return_intermediate='intensity')
# Normalize PSF by reference image
psf = image / norm
# Save PSF image to disk
if savepsfs:
filename_psf = f'psf_{zern_mode.name}_{zern_mode.convention + str(zern_mode.index)}_segs_{segment_pair[0]+1}-{segment_pair[1]+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_{zern_mode.name}_{zern_mode.convention + str(zern_mode.index)}_segs_{segment_pair[0]+1}-{segment_pair[1]+1}'
plt.clf()
hcipy.imshow_field(inter['seg_mirror'], grid=luv.aperture.grid, mask=luv.aperture, cmap='RdBu')
plt.savefig(os.path.join(resDir, 'OTE_images', opd_name + '.pdf'))
log.info('Calculating mean contrast in dark hole')
dh_intensity = psf * luv.dh_mask
contrast = np.mean(dh_intensity[np.where(luv.dh_mask != 0)])
log.info(f'contrast: {float(contrast)}') # contrast is a Field, here casting to normal float
return float(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 full_modes_from_themselves(instrument,
pmodes,
datadir,
sim_instance,
saving=False):
"""
Put all modes onto the segmented mirror in the pupil and get full 2D pastis modes, in pupil plane and focal plane.
Take the pmodes array of all modes (shape [segnum, modenum] = [nseg, nseg]) and apply them onto a segmented mirror
in the pupil. This phase gets returned both as an array of 2D arrays.
Both the pupl plane and the focal plane modes get save into a PDF grid, and as a cube to fits. Optionally, you can
save the pupil plane modes individually to PDF files by setting saving=True.
:param instrument: string, 'LUVOIR', 'HiCAT' or 'JWST'
:param pmodes: array of PASTIS modes [segnum, modenum], expected in nanometers
:param datadir: string, path to overall data directory containing matrix and results folder
:param sim_instance: class instance of the simulator for "instrument"
:param saving: bool, whether to save the individual pupil plane modes as PDFs to disk, default=False
:return: cube of pupil plane modes as array of 2D arrays
"""
nseg = pmodes.shape[0]
seglist = util.get_segment_list(instrument)
### Put all modes sequentially on the segmented mirror and get them as a phase map, then convert to WFE map
all_modes = []
all_modes_focal_plane = []
for i, thismode in enumerate(seglist):
if instrument == 'LUVOIR':
log.info(f'Working on mode {thismode}/{nseg}.')
wf_sm, wf_detector = util.apply_mode_to_luvoir(
pmodes[:, i], sim_instance)
psf_detector = wf_detector.intensity.shaped
all_modes_focal_plane.append(psf_detector)
all_modes.append(
(wf_sm.phase / wf_sm.wavenumber).shaped
) # wf_sm.phase is in rad, so this converts it to meters
if instrument == 'HiCAT':
log.info(f'Working on mode {thismode}/{nseg-1}.')
for segnum in range(nseg):
sim_instance.iris_dm.set_actuator(segnum,
pmodes[segnum, i] / 1e9, 0,
0) # /1e9 converts to meters
psf_detector_data, inter = sim_instance.calc_psf(
return_intermediates=True)
psf_detector = psf_detector_data[0].data
all_modes_focal_plane.append(psf_detector)
phase_sm = inter[1].phase
hicat_wavenumber = 2 * np.pi / (CONFIG_PASTIS.getfloat(
'HiCAT', 'lambda') / 1e9) # /1e9 converts to meters
all_modes.append(
phase_sm / hicat_wavenumber
) # phase_sm is in rad, so this converts it to meters
if instrument == 'JWST':
log.info(f'Working on mode {thismode}/{nseg - 1}.')
sim_instance[1].zero()
for segnum in range(
nseg
): # TODO: there is probably a single function that puts the aberration on the OTE at once
seg_name = webbpsf_imaging.WSS_SEGS[segnum].split('-')[0]
sim_instance[1].move_seg_local(seg_name,
piston=pmodes[segnum, i],
trans_unit='nm')
psf_detector_data, inter = sim_instance[0].calc_psf(
nlambda=1, return_intermediates=True)
psf_detector = psf_detector_data[0].data
all_modes_focal_plane.append(psf_detector)
phase_ote = inter[1].phase
jwst_wavenumber = 2 * np.pi / (CONFIG_PASTIS.getfloat(
'JWST', 'lambda') / 1e9) # /1e9 converts to meters
all_modes.append(
phase_ote / jwst_wavenumber
) # phase_sm is in rad, so this converts it to meters
### Check for results directory structure and create if it doesn't exist
log.info('Creating data directories')
subdirs = [
os.path.join(datadir, 'results'),
os.path.join(datadir, 'results', 'modes'),
os.path.join(datadir, 'results', 'modes', 'pupil_plane'),
os.path.join(datadir, 'results', 'modes', 'focal_plane'),
os.path.join(datadir, 'results', 'modes', 'focal_plane', 'fits'),
os.path.join(datadir, 'results', 'modes', 'pupil_plane', 'fits'),
os.path.join(datadir, 'results', 'modes', 'pupil_plane', 'pdf')
]
for place in subdirs:
if not os.path.isdir(place):
os.mkdir(place)
### Plot all modes together and save as PDF (pupil plane)
log.info('Saving all PASTIS modes together as PDF (pupil plane)...')
plt.figure(figsize=(36, 30))
for i, thismode in enumerate(seglist):
if instrument == 'LUVOIR':
plt.subplot(12, 10, i + 1)
if instrument == 'HiCAT':
plt.subplot(8, 5, i + 1)
if instrument == 'JWST':
plt.subplot(6, 3, i + 1)
plt.imshow(all_modes[i], cmap='RdBu')
plt.axis('off')
plt.title(f'Mode {thismode}')
plt.savefig(
os.path.join(datadir, 'results', 'modes', 'pupil_plane',
'modes_piston.pdf'))
### Plot them individually and save as PDF (pupil plane)
if saving:
log.info(
'Saving all PASTIS modes into individual PDFs (pupil plane)...')
for i, thismode in enumerate(seglist):
# pdf
plt.clf()
plt.imshow(
all_modes[i], cmap='RdBu'
) # TODO: this is now super slow for LUVOIR, using hcipy was way faster. Change back in LUVOIR case?
plt.axis('off')
plt.title(f'Mode {thismode}', size=30)
if saving:
plt.savefig(
os.path.join(datadir, 'results', 'modes', 'pupil_plane',
'pdf', f'mode_{thismode}.pdf'))
### Save as fits cube (pupil plane)
log.info('Saving all PASTIS modes into fits cube (pupil plane)')
mode_cube = np.array(all_modes)
hcipy.write_fits(
mode_cube,
os.path.join(datadir, 'results', 'modes', 'pupil_plane', 'fits',
'cube_modes.fits'))
### Plot all modes together and save as PDF (focal plane)
log.info('Saving all PASTIS modes together as PDF (focal plane)...')
plt.figure(figsize=(36, 30))
for i, thismode in enumerate(seglist):
if instrument == 'LUVOIR':
plt.subplot(12, 10, i + 1)
if instrument == 'HiCAT':
plt.subplot(8, 5, i + 1)
if instrument == 'JWST':
plt.subplot(6, 3, i + 1)
plt.imshow(all_modes_focal_plane[i], cmap='inferno', norm=LogNorm())
plt.axis('off')
plt.title(f'Mode {thismode}')
plt.savefig(
os.path.join(datadir, 'results', 'modes', 'focal_plane',
'modes_piston.pdf'))
### Save as fits cube (focal plane)
log.info('Saving all PASTIS modes into fits cube (focal plane)')
psf_cube = np.array(all_modes_focal_plane)
hcipy.write_fits(
psf_cube,
os.path.join(datadir, 'results', 'modes', 'focal_plane', 'fits',
'cube_modes.fits'))
return mode_cube
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 num_matrix_luvoir(design):
"""
Generate a numerical PASTIS matrix for a LUVOIR A coronagraph.
All inputs are read from the (local) configfile and saved to the specified output directory.
The LUVOIR STDT delivery in May 2018 included three different apodizers
we can work with, so I will implement an easy way of making a choice between them.
small, medium and large
"""
# Keep track of time
start_time = time.time() # runtime is currently around 150 minutes
print('Building numerical matrix for LUVOIR\n')
### Parameters
# System parameters
resDir = os.path.join(CONFIG_INI.get('local', 'local_data_path'), 'active',
'matrix_numerical')
zern_number = CONFIG_INI.getint('calibration', 'zernike')
zern_mode = util.ZernikeMode(
zern_number) # Create Zernike mode object for easier handling
# General telescope parameters
nb_seg = CONFIG_INI.getint('LUVOIR', 'nb_subapertures')
wvln = CONFIG_INI.getfloat('LUVOIR', 'lambda') * 1e-9 # m
diam = CONFIG_INI.getfloat('LUVOIR', 'diameter') # m
nm_aber = CONFIG_INI.getfloat('calibration',
'single_aberration') * 1e-9 # m
# Image system parameters
im_lamD = 30 # image size in lambda/D
sampling = 4
# Print some of the defined parameters
print('LUVOIR apodizer design: {}'.format(design))
print()
print('Wavelength: {} m'.format(wvln))
print('Telescope diameter: {} m'.format(diam))
print('Number of segments: {}'.format(nb_seg))
print()
print('Image size: {} lambda/D'.format(im_lamD))
print('Sampling: {} px per lambda/D'.format(sampling))
### Setting up the paths
# If subfolder "matrix_numerical" doesn't exist yet, create it.
if not os.path.isdir(resDir):
os.mkdir(resDir)
# If subfolder "OTE_images" doesn't exist yet, create it.
if not os.path.isdir(os.path.join(resDir, 'OTE_images')):
os.mkdir(os.path.join(resDir, 'OTE_images'))
# If subfolder "psfs" doesn't exist yet, create it.
if not os.path.isdir(os.path.join(resDir, 'psfs')):
os.mkdir(os.path.join(resDir, 'psfs'))
### Instantiate Luvoir telescope with chosen apodizer design
optics_input = '/Users/ilaginja/Documents/LabWork/ultra/LUVOIR_delivery_May2019/'
luvoir = LuvoirAPLC(optics_input, design, sampling)
### Dark hole mask
dh_outer = hc.circular_aperture(2 * luvoir.apod_dict[design]['owa'] *
luvoir.lam_over_d)(luvoir.focal_det)
dh_inner = hc.circular_aperture(2 * luvoir.apod_dict[design]['iwa'] *
luvoir.lam_over_d)(luvoir.focal_det)
dh_mask = (dh_outer - dh_inner).astype('bool')
### Reference images for contrast normalization and coronagraph floor
unaberrated_coro_psf, ref = luvoir.calc_psf(ref=True,
display_intermediate=False,
return_intermediate=False)
norm = np.max(ref)
dh_intensity = unaberrated_coro_psf / norm * dh_mask
contrast_floor = np.mean(dh_intensity[np.where(dh_intensity != 0)])
print(contrast_floor)
### Generating the PASTIS matrix and a list for all contrasts
matrix_direct = np.zeros([nb_seg, nb_seg]) # Generate empty matrix
all_psfs = []
all_contrasts = []
print('nm_aber: {} m'.format(nm_aber))
for i in range(nb_seg):
for j in range(nb_seg):
print('\nSTEP: {}-{} / {}-{}'.format(i + 1, j + 1, nb_seg, nb_seg))
# Put aberration on correct segments. If i=j, apply only once!
luvoir.flatten()
luvoir.set_segment(i + 1, nm_aber / 2, 0, 0)
if i != j:
luvoir.set_segment(j + 1, nm_aber / 2, 0, 0)
print('Calculating coro image...')
image, inter = luvoir.calc_psf(ref=False,
display_intermediate=False,
return_intermediate='intensity')
# Normalize PSF by reference image
psf = image / norm
# Save image to disk
filename_psf = 'psf_' + zern_mode.name + '_' + zern_mode.convention + str(
zern_mode.index) + '_segs_' + str(i + 1) + '-' + str(j + 1)
hc.write_fits(psf,
os.path.join(resDir, 'psfs', filename_psf + '.fits'))
all_psfs.append(psf)
# Save OPD images for testing (are these actually surface images, not OPD?)
opd_name = 'opd_' + zern_mode.name + '_' + zern_mode.convention + str(
zern_mode.index) + '_segs_' + str(i + 1) + '-' + str(j + 1)
plt.clf()
hc.imshow_field(inter['seg_mirror'],
mask=luvoir.aperture,
cmap='RdBu')
plt.savefig(os.path.join(resDir, 'OTE_images', opd_name + '.pdf'))
print('Calculating mean contrast in dark hole')
dh_intensity = psf * dh_mask
contrast = np.mean(dh_intensity[np.where(dh_intensity != 0)])
print('contrast:', contrast)
all_contrasts.append(contrast)
# Fill according entry in the matrix and subtract baseline contrast
matrix_direct[i, j] = contrast - contrast_floor
# Transform saved lists to arrays
all_psfs = np.array(all_psfs)
all_contrasts = np.array(all_contrasts)
# Filling the off-axis elements
matrix_two_N = np.copy(
matrix_direct
) # This is just an intermediary copy so that I don't mix things up.
matrix_pastis = np.copy(
matrix_direct) # 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
print('Off-axis for i{}-j{}: {}'.format(
i + 1, j + 1, matrix_off_val))
# Normalize matrix for the input aberration - the whole code is set up to be normalized to 1 nm, and even if
# the units entered are in m for the sake of HCIPy, everything else is assuming the baseline is 1nm, so the
# normalization can be taken out if we're working with exactly 1 nm for the aberration, even if entered in meters.
#matrix_pastis /= np.square(nm_aber)
# Save matrix to file
filename_matrix = 'PASTISmatrix_num_' + zern_mode.name + '_' + zern_mode.convention + str(
zern_mode.index)
hc.write_fits(matrix_pastis, os.path.join(resDir,
filename_matrix + '.fits'))
print('Matrix saved to:', os.path.join(resDir, filename_matrix + '.fits'))
# Save the PSF image *cube* as well (as opposed to each one individually)
hc.write_fits(
all_psfs,
os.path.join(resDir, 'psfs', 'psf_cube' + '.fits'),
)
np.savetxt(os.path.join(resDir, 'contrasts.txt'), all_contrasts, fmt='%e')
# Tell us how long it took to finish.
end_time = time.time()
print('Runtime for matrix_building.py:', end_time - start_time, 'sec =',
(end_time - start_time) / 60, 'min')
print('Data saved to {}'.format(resDir))
def full_modes_from_themselves(pmodes, datadir, sm, wf_aper, saving=False):
"""
Put all modes onto the pupuil SM and get full 2D modes.
Take the pmodes array of all modes (shape [segnum, modenum] = [nseg, nseg]) and apply them onto a segmenter mirror
in the pupil. This phase gets returned both as an array of hcipy.Fields, as well as a standard array of 2D arrays.
Optionally, save a PDF displaying all modes, a fits cube and individual PDF images.
:param pmodes: array of PASTIS modes [segnum, modenum]
:param datadir: string, path to overall data directory containing matrix and results folder
:param saving: bool, whether to save figure to disk or not
:return: all_modes as array of Fields, mode_cube as array of 2D arrays (hcipy vs matplotlib)
"""
nseg = 120
### Put all modes on the SM and get their phase
all_modes = []
for thismode in range(nseg):
print('Working on mode {}/{}.'.format(thismode + 1, nseg))
wf_sm = apply_mode_to_sm(pmodes[:, thismode], sm, wf_aper)
all_modes.append(wf_sm.phase / wf_sm.wavenumber
) # wf.phase is in rad, so this converts it to meters
### Check for results directory structure and create if it doesn't exist
if saving:
subdirs = [
os.path.join(datadir, 'results'),
os.path.join(datadir, 'results', 'modes'),
os.path.join(datadir, 'results', 'modes', 'fits'),
os.path.join(datadir, 'results', 'modes', 'pdf')
]
for place in subdirs:
if not os.path.isdir(place):
os.mkdir(place)
### Plot all modes together and save
if saving:
print('Saving all PASTIS modes...')
plt.figure(figsize=(36, 30))
for thismode in range(nseg):
plt.subplot(12, 10, thismode + 1)
hc.imshow_field(all_modes[thismode], cmap='RdBu')
plt.axis('off')
plt.title('Mode ' + str(thismode + 1))
plt.savefig(
os.path.join(datadir, 'results', 'modes', 'modes_piston.pdf'))
### Plot them individually and save as fits and pdf
mode_cube = [] # to save as a fits cube
for thismode in range(nseg):
# pdf
plt.clf()
hc.imshow_field(all_modes[thismode], cmap='RdBu')
plt.axis('off')
plt.title('Mode ' + str(thismode + 1), size=30)
if saving:
plt.savefig(
os.path.join(datadir, 'results', 'modes', 'pdf',
'mode' + str(thismode + 1) + '.pdf'))
# for the fits cube
mode_cube.append(all_modes[thismode].shaped)
# fits cube
mode_cube = np.array(mode_cube)
if saving:
hc.write_fits(
mode_cube,
os.path.join(datadir, 'results', 'modes', 'fits',
'cube_modes.fits'))
return all_modes, mode_cube
def num_matrix_multiprocess(instrument, design=None, savepsfs=True, saveopds=True):
"""
Generate a numerical/semi-analytical PASTIS matrix.
Multiprocessed script to calculate PASTIS matrix. Implementation adapted from
hicat.scripts.stroke_minimization.calculate_jacobian
:param instrument: str, what instrument (LUVOIR, HiCAT, JWST) to generate the PASTIS matrix for
:param design: str, optional, default=None, which means we read from the configfile: what coronagraph design
to use - 'small', 'medium' or 'large'
: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
:return: overall_dir: string, experiment directory
"""
# Keep track of time
start_time = time.time() # runtime is currently around 150 minutes
### Parameters
# Create directory names
tel_suffix = f'{instrument.lower()}'
if instrument == 'LUVOIR':
if design is None:
design = CONFIG_PASTIS.get('LUVOIR', 'coronagraph_design')
tel_suffix += f'-{design}'
overall_dir = util.create_data_path(CONFIG_PASTIS.get('local', 'local_data_path'), telescope=tel_suffix)
os.makedirs(overall_dir, exist_ok=True)
resDir = os.path.join(overall_dir, 'matrix_numerical')
# Create necessary directories if they don't exist yet
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)
# Set up logger
util.setup_pastis_logging(resDir, f'pastis_matrix_{tel_suffix}')
log.info(f'Building numerical matrix for {tel_suffix}\n')
# Read calibration aberration
zern_number = CONFIG_PASTIS.getint('calibration', 'local_zernike')
zern_mode = util.ZernikeMode(zern_number) # Create Zernike mode object for easier handling
# General telescope parameters
nb_seg = CONFIG_PASTIS.getint(instrument, 'nb_subapertures')
seglist = util.get_segment_list(instrument)
wvln = CONFIG_PASTIS.getfloat(instrument, 'lambda') * 1e-9 # m
wfe_aber = CONFIG_PASTIS.getfloat(instrument, 'calibration_aberration') * 1e-9 # m
# Record some of the defined parameters
log.info(f'Instrument: {tel_suffix}')
log.info(f'Wavelength: {wvln} m')
log.info(f'Number of segments: {nb_seg}')
log.info(f'Segment list: {seglist}')
log.info(f'wfe_aber: {wfe_aber} m')
log.info(f'Total number of segment pairs in {instrument} pupil: {len(list(util.segment_pairs_all(nb_seg)))}')
log.info(f'Non-repeating pairs in {instrument} pupil calculated here: {len(list(util.segment_pairs_non_repeating(nb_seg)))}')
# Copy configfile to resulting matrix directory
util.copy_config(resDir)
# Calculate coronagraph floor, and normalization factor from direct image
contrast_floor, norm = calculate_unaberrated_contrast_and_normalization(instrument, design, return_coro_simulator=False,
save_coro_floor=True, save_psfs=False, outpath=overall_dir)
# Figure out how many processes is optimal and create a Pool.
# Assume we're the only one on the machine so we can hog all the resources.
# We expect numpy to use multithreaded math via the Intel MKL library, so
# we check how many threads MKL will use, and create enough processes so
# as to use 100% of the CPU cores.
# You might think we should divide number of cores by 2 to get physical cores
# to account for hyperthreading, however empirical testing on telserv3 shows that
# it is slightly more performant on telserv3 to use all logical cores
num_cpu = multiprocessing.cpu_count()
# try:
# import mkl
# num_core_per_process = mkl.get_max_threads()
# except ImportError:
# # typically this is 4, so use that as default
# log.info("Couldn't import MKL; guessing default value of 4 cores per process")
# num_core_per_process = 4
num_core_per_process = 1 # NOTE: this was changed by Scott Will in HiCAT and makes more sense, somehow
num_processes = int(num_cpu // num_core_per_process)
log.info(f"Multiprocess PASTIS matrix for {instrument} will use {num_processes} processes (with {num_core_per_process} threads per process)")
# Set up a function with all arguments fixed except for the last one, which is the segment pair tuple
if instrument == 'LUVOIR':
calculate_matrix_pair = functools.partial(_luvoir_matrix_one_pair, design, norm, wfe_aber, zern_mode, resDir,
savepsfs, saveopds)
if instrument == 'HiCAT':
# Copy used BostonDM maps to matrix folder
shutil.copytree(CONFIG_PASTIS.get('HiCAT', 'dm_maps_path'), os.path.join(resDir, 'hicat_boston_dm_commands'))
calculate_matrix_pair = functools.partial(_hicat_matrix_one_pair, norm, wfe_aber, resDir, savepsfs, saveopds)
if instrument == 'JWST':
calculate_matrix_pair = functools.partial(_jwst_matrix_one_pair, norm, wfe_aber, resDir, savepsfs, saveopds)
# Iterate over all segment pairs via a multiprocess pool
mypool = multiprocessing.Pool(num_processes)
t_start = time.time()
results = mypool.map(calculate_matrix_pair, util.segment_pairs_non_repeating(nb_seg)) # this util function returns a generator
t_stop = time.time()
log.info(f"Multiprocess calculation complete in {t_stop-t_start}sec = {(t_stop-t_start)/60}min")
# Unscramble results
# results is a list of tuples that contain the return from the partial function, in this case: result[i] = (c, (seg1, seg2))
contrast_matrix = np.zeros([nb_seg, nb_seg]) # Generate empty matrix
for i in range(len(results)):
# Fill according entry in the matrix and subtract baseline contrast
contrast_matrix[results[i][1][0], results[i][1][1]] = results[i][0] - contrast_floor
mypool.close()
# Save all contrasts to disk, WITH subtraction of coronagraph floor
hcipy.write_fits(contrast_matrix, os.path.join(resDir, 'pair-wise_contrasts.fits'))
plt.figure(figsize=(10, 10))
plt.imshow(contrast_matrix)
plt.colorbar()
plt.savefig(os.path.join(resDir, 'contrast_matrix.pdf'))
# Calculate the PASTIS matrix from the contrast matrix: off-axis elements and normalization
matrix_pastis = pastis_from_contrast_matrix(contrast_matrix, seglist, wfe_aber)
# Save matrix to file
filename_matrix = f'PASTISmatrix_num_{zern_mode.name}_{zern_mode.convention + str(zern_mode.index)}'
hcipy.write_fits(matrix_pastis, os.path.join(resDir, filename_matrix + '.fits'))
ppl.plot_pastis_matrix(matrix_pastis, wvln*1e9, out_dir=resDir, save=True) # convert wavelength to nm
log.info(f'Matrix saved to: {os.path.join(resDir, filename_matrix + ".fits")}')
# Tell us how long it took to finish.
end_time = time.time()
log.info(f'Runtime for matrix_building_numerical.py/multiprocess: {end_time - start_time}sec = {(end_time - start_time)/60}min')
log.info(f'Data saved to {resDir}')
return overall_dir
def num_matrix_luvoir(design, savepsfs=False, saveopds=True):
"""
Generate a numerical PASTIS matrix for a LUVOIR A 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.
The LUVOIR STDT delivery in May 2018 included three different apodizers
we can work with, you pick which of the three you want with the 'design' parameter.
:param design: string, what coronagraph design to use - 'small', 'medium' or 'large'
:param savepsfs: bool, if True, all PSFs will be saved to disk individually, as fits files, additionally to the
total PSF cube. If False, the total cube will still get saved at the very end of the script.
:param saveopds: bool, if True, all pupil surface maps of aberrated segment pairs will be saved to disk as PDF
:return overall_dir: string, experiment directory
"""
# Keep track of time
start_time = time.time()
### Parameters
# System parameters
overall_dir = util.create_data_path(CONFIG_PASTIS.get('local', 'local_data_path'), telescope='luvoir-'+design)
os.makedirs(overall_dir, exist_ok=True)
resDir = os.path.join(overall_dir, 'matrix_numerical')
# Create necessary directories if they don't exist yet
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)
# Set up logger
util.setup_pastis_logging(resDir, f'pastis_matrix_{design}')
log.info('Building numerical matrix for LUVOIR\n')
# Read calibration aberration
zern_number = CONFIG_PASTIS.getint('calibration', 'local_zernike')
zern_mode = util.ZernikeMode(zern_number) # Create Zernike mode object for easier handling
# General telescope parameters
nb_seg = CONFIG_PASTIS.getint('LUVOIR', 'nb_subapertures')
wvln = CONFIG_PASTIS.getfloat('LUVOIR', 'lambda') * 1e-9 # m
diam = CONFIG_PASTIS.getfloat('LUVOIR', 'diameter') # m
wfe_aber = CONFIG_PASTIS.getfloat('LUVOIR', 'calibration_aberration') * 1e-9 # m
# Image system parameters
sampling = CONFIG_PASTIS.getfloat('LUVOIR', 'sampling')
# Record some of the defined parameters
log.info(f'LUVOIR apodizer design: {design}')
log.info(f'Wavelength: {wvln} m')
log.info(f'Telescope diameter: {diam} m')
log.info(f'Number of segments: {nb_seg}')
log.info(f'Sampling: {sampling} px per lambda/D')
log.info(f'wfe_aber: {wfe_aber} m')
# Copy configfile to resulting matrix directory
util.copy_config(resDir)
### Instantiate Luvoir telescope with chosen apodizer design
optics_input = CONFIG_PASTIS.get('LUVOIR', 'optics_path')
luvoir = LuvoirAPLC(optics_input, design, sampling)
### Reference images for contrast normalization and coronagraph floor
unaberrated_coro_psf, ref = luvoir.calc_psf(ref=True, display_intermediate=False, return_intermediate=False)
norm = np.max(ref)
dh_intensity = (unaberrated_coro_psf / norm) * luvoir.dh_mask
contrast_floor = np.mean(dh_intensity[np.where(luvoir.dh_mask != 0)])
log.info(f'contrast floor: {contrast_floor}')
### Generating the PASTIS matrix and a list for all contrasts
contrast_matrix = np.zeros([nb_seg, nb_seg]) # Generate empty matrix
all_psfs = []
all_contrasts = []
for i in range(nb_seg):
for j in range(nb_seg):
log.info(f'\nSTEP: {i+1}-{j+1} / {nb_seg}-{nb_seg}')
# Put aberration on correct segments. If i=j, apply only once!
luvoir.flatten()
luvoir.set_segment(i+1, wfe_aber/2, 0, 0)
if i != j:
luvoir.set_segment(j+1, wfe_aber/2, 0, 0)
log.info('Calculating coro image...')
image, inter = luvoir.calc_psf(ref=False, display_intermediate=False, return_intermediate='intensity')
# Normalize PSF by reference image
psf = image / norm
all_psfs.append(psf.shaped)
# Save image to disk
if savepsfs: # TODO: I might want to change this to matplotlib images since I save the PSF cube anyway.
filename_psf = f'psf_{zern_mode.name}_{zern_mode.convention + str(zern_mode.index)}_segs_{i+1}-{j+1}'
hcipy.write_fits(psf, os.path.join(resDir, 'psfs', filename_psf + '.fits'))
# Save OPD images for testing
if saveopds:
opd_name = f'opd_{zern_mode.name}_{zern_mode.convention + str(zern_mode.index)}_segs_{i+1}-{j+1}'
plt.clf()
hcipy.imshow_field(inter['seg_mirror'], mask=luvoir.aperture, cmap='RdBu')
plt.savefig(os.path.join(resDir, 'OTE_images', opd_name + '.pdf'))
log.info('Calculating mean contrast in dark hole')
dh_intensity = psf * luvoir.dh_mask
contrast = np.mean(dh_intensity[np.where(luvoir.dh_mask != 0)])
log.info(f'contrast: {float(contrast)}') # contrast is a Field, here casting to normal float
all_contrasts.append(contrast)
# Fill according entry in the matrix and subtract baseline contrast
contrast_matrix[i,j] = contrast - contrast_floor
# Transform saved lists to arrays
all_psfs = np.array(all_psfs)
all_contrasts = np.array(all_contrasts)
# Save the PSF image *cube* as well (as opposed to each one individually)
hcipy.write_fits(all_psfs, os.path.join(resDir, 'psfs', 'psf_cube.fits'),)
np.savetxt(os.path.join(resDir, 'pair-wise_contrasts.txt'), all_contrasts, fmt='%e')
# Filling the off-axis elements
log.info('\nCalculating off-axis matrix 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 - this defines what units the PASTIS matrix will be in. The PASTIS
# matrix propagation function (util.pastis_contrast()) then needs to take in the aberration vector in these same
# units. I have chosen to keep this to 1nm, so, we normalize the PASTIS matrix to units of nanometers.
matrix_pastis /= np.square(wfe_aber * 1e9) # 1e9 converts the calibration aberration back to nanometers
# Save matrix to file
filename_matrix = f'PASTISmatrix_num_{zern_mode.name}_{zern_mode.convention + str(zern_mode.index)}'
hcipy.write_fits(matrix_pastis, os.path.join(resDir, filename_matrix + '.fits'))
log.info(f'Matrix saved to: {os.path.join(resDir, filename_matrix + ".fits")}')
# 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}')
return overall_dir